Added Gemini support and refactored code

app.py

@@ -1,140 +1,208 @@
-import gradio as gr
 import base64
-from io import BytesIO
+import json
+import io
 import re
+from io import BytesIO
 from PIL import Image
+import gradio as gr
 from openai import OpenAI
+from google import genai
+from google.genai import types
+
+# ------------------ #
+# Utility Functions
+# ------------------ #
 
-# Helper to convert image to base64
 def image_to_base64(image: Image.Image) -> str:
+    """Convert a PIL image to base64 string."""
     buffered = BytesIO()
     image.save(buffered, format="PNG")
     return base64.b64encode(buffered.getvalue()).decode()
 
-# 🧠 Main logic function
-def recyclo_advisor(item_description, item_image, location, api_key):
-    if not api_key or api_key.strip() == "":
-        return "❗ Please enter a valid Nebius API key.", "", "", ""
-
+def parse_nebius_response(content: str):
+    """Parse the response content from Nebius."""
+    item_name = re.search(r"Item Name:\s*(.*)", content)
+    status = re.search(r"1\.\s*Recyclability Status:\s*(.*)", content)
+    instructions = re.search(r"2\.\s*Instructions:\s*(.*?)(?:3\.|$)", content, re.DOTALL)
+    reasoning = re.search(r"3\.\s*Reasoning:\s*(.*)", content, re.DOTALL)
+
+    return {
+        "item_name": item_name.group(1).strip() if item_name else "This item",
+        "status": status.group(1).strip() if status else "Unknown",
+        "instructions": instructions.group(1).strip() if instructions else "",
+        "reasoning": reasoning.group(1).strip() if reasoning else ""
+    }
+
+def parse_gemini_response(response_text: str):
+    """Parse the JSON string response from Gemini provider."""
     try:
-        client = OpenAI(
-            base_url="https://api.studio.nebius.com/v1/",
-            api_key=api_key
-        )
-    except Exception as e:
-        return "❌ API key error", "", "", str(e)
+        data = json.loads(response_text)
+    except json.JSONDecodeError:
+        data = {}
+    return {
+        "item_name": data.get("Item Name", "This item").strip(),
+        "status": data.get("1. Recyclability Status", "").strip(),
+        "instructions": data.get("2. Instructions", "").strip(),
+        "reasoning": data.get("3. Reasoning", "").strip()
+    }
+
+def build_prompt(item_description: str, location: str):
+    """Build system and user prompts."""
+    system_prompt = (
+        "You are a recycling and waste management expert. "
+        "Your job is to help users determine whether an item is recyclable, and if not, guide them on responsible disposal based on their location. "
+        "Be specific, practical, and locally relevant.\n\n"
+        "Always format your response as follows:\n"
+        "Item Name: <clear name>\n"
+        "1. Recyclability Status: Recyclable / Not Recyclable / Depends\n"
+        "2. Instructions: What should the user do with the item?\n"
+        "3. Reasoning: Why is this the right action in the selected location?"
+    )
 
-    try:
-        messages = [
-            {
-                "role": "system",
-                "content": (
-                    "You are a recycling and waste management expert. "
-                    "Your job is to help users determine whether an item is recyclable, and if not, guide them on responsible disposal based on their location. "
-                    "Be specific, practical, and locally relevant.\n\n"
-                    "Always format your response as follows:\n"
-                    "Item Name: <clear name>\n"
-                    "1. ♻️ Recyclability Status: Recyclable / Not Recyclable / Depends\n"
-                    "2. 🧾 Instructions: What should the user do with the item?\n"
-                    "3. ℹ️ Reasoning: Why is this the right action in the selected location?"
-                )
-            }
-        ]
+    description = f"The user is located in: {location}.\n"
+    description += f'They described the item as: "{item_description}".' if item_description else "No description provided."
 
-        user_content = []
-        description_text = f"The user is located in: {location}.\n"
-        if item_description:
-            description_text += f'They described the item as: "{item_description}".'
-        else:
-            description_text += "No description provided."
+    return system_prompt, description
+
+def validate_inputs(api_key, item_image, location):
+    """Validate required inputs."""
+    if not api_key or api_key.strip() == "":
+        raise gr.Error("🔐 API Key is required.")
+    if not item_image:
+        raise gr.Error("📷 Please upload image of the item.")
+    if not location or location.strip() == "":
+        raise gr.Error("🌍 Please enter your region.")
+
+# ------------------------- #
+# Main Processing Function
+# ------------------------- #
+
+def recyclo_advisor(item_description, item_image, location, api_key, provider):
+    """Main advisor logic: processes the image and description using the chosen provider."""
+    validate_inputs(api_key, item_image, location)
 
-        user_content.append({"type": "text", "text": description_text})
+    try:
+        system_prompt, user_text = build_prompt(item_description, location)
+        user_content = [{"type": "text", "text": user_text}]
 
         if item_image:
-            image_b64 = image_to_base64(item_image)
+            img_b64 = image_to_base64(item_image)
             user_content.append({
                 "type": "image_url",
-                "image_url": {"url": f"data:image/png;base64,{image_b64}"}
+                "image_url": {"url": f"data:image/png;base64,{img_b64}"}
             })
 
-        messages.append({"role": "user", "content": user_content})
-
-        response = client.chat.completions.create(
-            model="google/gemma-3-27b-it",
-            messages=messages,
-            max_tokens=2048,
-            temperature=0.6,
-            top_p=0.9
-        )
-
-        full_response = response.choices[0].message.content.strip()
-
-        # Extract Item Name, Status, Instructions, Reasoning
-        item_match = re.search(r"Item Name:\s*(.*)", full_response)
-        item_name = item_match.group(1).strip() if item_match else "This item"
-
-        status_match = re.search(r"1\.\s*♻️ Recyclability Status:\s*(.*)", full_response)
-        status_raw = status_match.group(1).strip() if status_match else "Unknown"
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": user_content}
+        ]
 
-        instructions_match = re.search(r"2\.\s*🧾 Instructions:\s*(.*?)(?:3\.|$)", full_response, re.DOTALL)
-        instructions = instructions_match.group(1).strip() if instructions_match else ""
+        if provider == "Nebius":
+            client = OpenAI(
+                base_url="https://api.studio.nebius.com/v1/",
+                api_key=api_key
+            )
+            response = client.chat.completions.create(
+                model="google/gemma-3-27b-it",
+                messages=messages,
+                max_tokens=2048,
+                temperature=0.6,
+                top_p=0.9
+            )
+            full_response = response.choices[0].message.content.strip()
+            result = parse_nebius_response(full_response)
+
+        else:  # Gemini
+            client = genai.Client(api_key=api_key)
+            prompt = system_prompt + "\n" + user_text
+
+            image_obj = None
+            for part in user_content:
+                if part["type"] == "image_url":
+                    b64_data = part["image_url"]["url"].split(",")[1]
+                    image_bytes = base64.b64decode(b64_data)
+                    image_obj = Image.open(io.BytesIO(image_bytes))
+                    break
+
+            if not image_obj:
+                raise ValueError("No image provided for Gemini provider")
+
+            config = types.GenerateContentConfig(response_mime_type="application/json")
+            response = client.models.generate_content(
+                model="gemini-2.5-flash",
+                contents=[image_obj, prompt],
+                config=config
+            )
+            result = parse_gemini_response(response.text.strip())
 
-        reasoning_match = re.search(r"3\.\s*ℹ️ Reasoning:\s*(.*)", full_response, re.DOTALL)
-        reasoning = reasoning_match.group(1).strip() if reasoning_match else ""
+        # Label formatting
+        status = result["status"]
+        item_name = result["item_name"]
 
-        # Label summary with item name included
-        if "Not Recyclable" in status_raw:
+        if "Not Recyclable" in status:
             label = f"❌ {item_name} is Not Recyclable"
-        elif "Depends" in status_raw:
+        elif "Depends" in status:
             label = f"⚠️ {item_name} recyclability Depends"
         else:
             label = f"✅ {item_name} is Recyclable"
 
-        # Basic summary report
         markdown_summary = f"""
         ### ♻️ **Recyclability Report for `{item_name}`**
         **Recyclability Status:**  
-        {status_raw}
+        {status}
         """
 
-        return label, markdown_summary, instructions, reasoning
+        return label, markdown_summary.strip(), result["instructions"], result["reasoning"]
 
     except Exception as e:
-        return "❌ Error calling Nebius LLM", "", "", str(e)
+        print(f"[Error] {e}")
+        return "❌ Error calling API Endpoint", "", "", str(e)
+
+# ----------- #
+# Gradio UI
+# ----------- #
 
-# 🎨 Gradio UI
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("## ♻️ RecycloBot – AI-Powered Recycling Advisor")
     gr.Markdown("Upload an item photo or describe it to get location-aware recycling guidance.")
 
     with gr.Row():
         with gr.Column(scale=1):
+            provider_dropdown = gr.Dropdown(
+                choices=["Nebius", "Gemini"],
+                value="Nebius",
+                label="🧠 Select Provider"
+            )
+
             api_key = gr.Textbox(
-                label="🔐 Nebius API Key",
-                placeholder="Paste your Nebius API key here",
+                label="🔐 API Key (Nebius)",
+                placeholder="Paste your API key here",
                 type="password"
             )
 
+            def update_label(provider):
+                return gr.update(label=f"🔐 API Key ({provider})")
+
+            provider_dropdown.change(fn=update_label, inputs=provider_dropdown, outputs=api_key)
+
             item_image = gr.Image(label="📷 Upload Image (Optional)", type="pil")
-            
+
             item_description = gr.Textbox(
-                label="📝 Describe the Item",
+                label="📝 Describe the Item (OPTIONAL)",
                 placeholder="e.g., 'USB cable', 'Greasy pizza box'"
             )
 
             location = gr.Textbox(
+                label="🌍 Your Region",
                 placeholder="Please input your location. e.g., Country and/or state name",
-                value="USA",
-                label="🌍 Your Region"
+                value="USA"
             )
 
             submit_btn = gr.Button("🚀 Analyze Item")
-        
+
             examples = gr.Examples(
-                examples=[
-                    ["broken phone", "img/broken_phone.jpg", "India"],
-                    ["charger", "img/charger.jpg", "London, England"]
-                ],
+                examples=[["broken phone", "img/broken_phone.jpg", "India"],
+                          ["charger", "img/charger.jpg", "London, England"]],
                 inputs=[item_description, item_image, location],
                 label="🧪 Try an Example"
             )
@@ -142,14 +210,14 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         with gr.Column(scale=1):
             status_output = gr.Label(label="♻️ Item Recyclability Summary")
             summary_output = gr.Markdown(label="📋 Recyclability Report")
-            with gr.Accordion("📦 What To Do With It", open=False):
+            with gr.Accordion("ℹ️ Instructions on What Exactly to Do with It", open=False):
                 instructions_output = gr.Markdown()
-            with gr.Accordion("ℹ️ Why This Matters", open=False):
+            with gr.Accordion("📦 Why This Matters", open=False):
                 reasoning_output = gr.Markdown()
 
     submit_btn.click(
         fn=recyclo_advisor,
-        inputs=[item_description, item_image, location, api_key],
+        inputs=[item_description, item_image, location, api_key, provider_dropdown],
         outputs=[status_output, summary_output, instructions_output, reasoning_output]
     )
 

requirements.txt

@@ -1,3 +1,4 @@
 gradio
 pillow
-openai
+openai
+google-genai