Update core/gemini_handler.py

core/gemini_handler.py

@@ -1,7 +1,7 @@
 # core/gemini_handler.py
 import google.generativeai as genai
 import json
-import re # For cleaning JSON output
+import re
 
 class GeminiHandler:
     def __init__(self, api_key):
@@ -9,29 +9,57 @@ class GeminiHandler:
         self.model = genai.GenerativeModel('gemini-1.5-flash-latest') # Or your preferred Gemini model
 
     def _clean_json_response(self, text_response):
-        # Gemini might sometimes wrap JSON in ```json ... ```
-        match = re.search(r"```json\s*([\s\S]*?)\s*```", text_response)
+        # Attempt to find JSON within backticks or directly
+        match = re.search(r"```json\s*([\s\S]*?)\s*```", text_response, re.DOTALL)
         if match:
-            return match.group(1).strip()
-        return text_response.strip()
+            json_str = match.group(1).strip()
+        else:
+            # Try to find the start of a list or object
+            json_start_list = text_response.find('[')
+            json_start_obj = text_response.find('{')
 
-    def generate_story_breakdown(self, prompt_text):
-        try:
-            response = self.model.generate_content(prompt_text)
-            cleaned_response = self._clean_json_response(response.text)
-            # print(f"Debug: Cleaned Story Breakdown JSON: {cleaned_response}") # For debugging
-            story_scenes = json.loads(cleaned_response)
-            return story_scenes
-        except Exception as e:
-            print(f"Error generating story breakdown: {e}")
-            print(f"Problematic Gemini Response: {response.text if 'response' in locals() else 'No response object'}")
-            return None
+            if json_start_list != -1 and (json_start_obj == -1 or json_start_list < json_start_obj):
+                json_str = text_response[json_start_list:]
+            elif json_start_obj != -1:
+                json_str = text_response[json_start_obj:]
+            else:
+                return text_response # Not clearly JSON
+
+        # Remove trailing characters that might break parsing if JSON is incomplete
+        # This is a bit aggressive, might need refinement
+        # Find last '}' or ']'
+        last_bracket = max(json_str.rfind('}'), json_str.rfind(']'))
+        if last_bracket != -1:
+            json_str = json_str[:last_bracket+1]
+        return json_str.strip()
 
-    def generate_image_prompt(self, prompt_text):
+    def _execute_gemini_call(self, prompt_text, expect_json=False):
         try:
             response = self.model.generate_content(prompt_text)
-            # Image prompts are usually just text, no need for JSON cleaning unless specified
-            return response.text.strip()
+            text_content = response.text
+            if expect_json:
+                cleaned_response = self._clean_json_response(text_content)
+                # print(f"DEBUG: Cleaned JSON for prompt:\n{prompt_text[:200]}...\nResponse:\n{cleaned_response}") # Debug
+                return json.loads(cleaned_response)
+            return text_content.strip()
+        except json.JSONDecodeError as e:
+            print(f"Error decoding JSON from Gemini response: {e}")
+            print(f"Problematic Gemini Raw Response:\n{text_content if 'text_content' in locals() else 'No response object'}")
+            print(f"Cleaned attempt was:\n{cleaned_response if 'cleaned_response' in locals() else 'N/A'}")
+            raise # Re-raise to be caught by caller
         except Exception as e:
-            print(f"Error generating image prompt: {e}")
-            return None
+            print(f"Error in Gemini call: {e}")
+            print(f"Problematic Gemini Raw Response (if available):\n{response.text if 'response' in locals() else 'No response object'}")
+            raise # Re-raise
+
+    def generate_story_breakdown(self, prompt_text):
+        return self._execute_gemini_call(prompt_text, expect_json=True)
+
+    def generate_image_prompt(self, prompt_text): # This is for generating a new image prompt string
+        return self._execute_gemini_call(prompt_text, expect_json=False)
+
+    def regenerate_scene_script_details(self, prompt_text): # Expects JSON for a single scene
+        return self._execute_gemini_call(prompt_text, expect_json=True)
+
+    def regenerate_image_prompt_from_feedback(self, prompt_text): # Expects a string (the new prompt)
+        return self._execute_gemini_call(prompt_text, expect_json=False)