Update app.py

app.py

@@ -15,6 +15,24 @@ except KeyError:
     raise ValueError("Please set the GEMINI_API_KEY environment variable.")
 client = genai.Client(api_key=api_key)
 
+def clean_response_text(response_text):
+    """
+    Clean the API response by removing Markdown code block markers.
+    
+    Args:
+        response_text (str): The raw response text from the API.
+    
+    Returns:
+        str: The cleaned response text.
+    """
+    # Remove ```json and ``` markers
+    cleaned_text = response_text.strip()
+    if cleaned_text.startswith("```json"):
+        cleaned_text = cleaned_text[len("```json"):].strip()
+    if cleaned_text.endswith("```"):
+        cleaned_text = cleaned_text[:-len("```")].strip()
+    return cleaned_text
+
 def generate_ideas(tag):
     """
     Generate a diverse set of ideas related to the tag using the LLM.
@@ -34,14 +52,16 @@ def generate_ideas(tag):
     """
     try:
         response = client.models.generate_content(
-            model='gemini-2.0-flash',  # Updated to a potentially supported model
+            model='gemini-2.0-flash',
             contents=[prompt],
             config=types.GenerateContentConfig(temperature=1.2)
         )
         print(f"Raw response for ideas: {response.text}")  # Debugging
         if not response.text or response.text.isspace():
             raise ValueError("Empty response from API")
-        response_json = json.loads(response.text.strip())
+        # Clean the response text
+        cleaned_text = clean_response_text(response.text)
+        response_json = json.loads(cleaned_text)
         if 'ideas' not in response_json or not isinstance(response_json['ideas'], list):
             raise ValueError("Invalid JSON format: 'ideas' key missing or not a list")
         return response_json['ideas']
@@ -78,14 +98,16 @@ def generate_item(tag, ideas):
     """
     try:
         response = client.models.generate_content(
-            model='gemini-2.0-flash',  # Updated to a potentially supported model
+            model='gemini-2.0-flash',
             contents=[prompt],
             config=types.GenerateContentConfig(temperature=1.2)
         )
         print(f"Raw response for item: {response.text}")  # Debugging
         if not response.text or response.text.isspace():
             raise ValueError("Empty response from API")
-        response_json = json.loads(response.text.strip())
+        # Clean the response text
+        cleaned_text = clean_response_text(response.text)
+        response_json = json.loads(cleaned_text)
         if 'caption' not in response_json or 'image_prompt' not in response_json:
             raise ValueError("Invalid JSON format: 'caption' or 'image_prompt' key missing")
         text = response_json['caption']
@@ -486,5 +508,5 @@ with gr.Blocks(
         outputs=[current_tag, current_index, feed_items, feed_html, is_loading]
     )
 
-# Launch the app with a public link
-demo.launch(share=True)
+# Launch the app (removed share=True since we're on Hugging Face Spaces)
+demo.launch()