File size: 3,197 Bytes
db99ae5 4afbb07 db99ae5 4afbb07 db99ae5 4afbb07 db99ae5 4afbb07 db99ae5 4afbb07 db99ae5 4afbb07 db99ae5 4afbb07 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 |
# core/gemini_handler.py
import google.generativeai as genai
import json
import re
class GeminiHandler:
def __init__(self, api_key):
genai.configure(api_key=api_key)
self.model = genai.GenerativeModel('gemini-1.5-flash-latest') # Or your preferred Gemini model
def _clean_json_response(self, 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:
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('{')
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 _execute_gemini_call(self, prompt_text, expect_json=False):
try:
response = self.model.generate_content(prompt_text)
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 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) |