Update sozo_gen.py

sozo_gen.py

@@ -24,6 +24,8 @@ from typing import Dict, List, Tuple, Any
 from langchain_google_genai import ChatGoogleGenerativeAI
 from google import genai
 import requests
+# In sozo_gen.py, near the other google imports
+from google.genai import types as genai_types
 
 # --- Configuration ---
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - [%(funcName)s] - %(message)s')
@@ -83,6 +85,122 @@ def clean_narration(txt: str) -> str:
 
 def placeholder_img() -> Image.Image: return Image.new("RGB", (WIDTH, HEIGHT), (230, 230, 230))
 
+# In sozo_gen.py, add these new functions at the end of the file
+
+def generate_image_with_gemini(prompt: str) -> Image.Image:
+    """Generates an image using the specified Gemini model and client configuration."""
+    logging.info(f"Generating Gemini image with prompt: '{prompt}'")
+    try:
+        # Use the genai.Client as per the correct implementation
+        client = genai.Client(api_key=API_KEY)
+        full_prompt = f"A professional, 3d digital art style illustration for a business presentation: {prompt}"
+        
+        response = client.models.generate_content(
+            model="gemini-2.0-flash-exp",
+            contents=full_prompt,
+            config=genai_types.GenerateContentConfig(
+                response_modalities=["Text", "Image"]
+            ),
+        )
+        
+        # Find the image part in the response
+        img_part = next((part for part in response.candidates[0].content.parts if part.content_type == "Image"), None)
+        
+        if img_part:
+            # The content is already bytes, so we can open it directly
+            return Image.open(io.BytesIO(img_part.content)).convert("RGB")
+        else:
+            logging.error("Gemini response did not contain an image.")
+            return None
+    except Exception as e:
+        logging.error(f"Gemini image generation failed: {e}")
+        return None
+
+def generate_slides_from_report(raw_md: str, chart_urls: dict, uid: str, project_id: str, bucket, llm):
+    """
+    Uses an AI planner to convert a report into a 10-slide presentation deck.
+    """
+    logging.info(f"Generating slides for project {project_id}")
+
+    planner_prompt = f"""
+    You are an expert presentation designer. Your task is to convert the following data analysis report into a concise and visually engaging 10-slide deck.
+
+    **Full Report Content:**
+    ---
+    {raw_md}
+    ---
+
+    **Instructions:**
+    1.  Read the entire report to understand the core narrative and key findings.
+    2.  Create a plan for exactly 10 slides.
+    3.  For each slide, define a `title` and short `content` (2-3 bullet points or a brief paragraph).
+    4.  For the visual on each slide, you must decide between two types:
+        - If a report section is supported by an existing chart (indicated by a `<generate_chart:...>` tag), you **must** use it. Set `visual_type: "existing_chart"` and `visual_ref: "the exact chart description from the tag"`.
+        - For key points without a chart (like introductions, conclusions, or text-only insights), you **must** request a new image. Set `visual_type: "new_image"` and `visual_ref: "a concise, descriptive prompt for an AI to generate a 3D digital art style illustration"`.
+    5.  You must request exactly 3-4 new images to balance the presentation.
+
+    **Output Format:**
+    Return ONLY a valid JSON array of 10 slide objects. Do not include any other text or markdown formatting.
+
+    Example:
+    [
+      {{ "slide_number": 1, "title": "Introduction", "content": "...", "visual_type": "new_image", "visual_ref": "A 3D illustration of a rising stock chart" }},
+      {{ "slide_number": 2, "title": "Sales by Region", "content": "...", "visual_type": "existing_chart", "visual_ref": "bar | Sales by Region" }},
+      ...
+    ]
+    """
+
+    try:
+        plan_response = llm.invoke(planner_prompt).content.strip()
+        if plan_response.startswith("```json"):
+            plan_response = plan_response[7:-3]
+        slide_plan = json.loads(plan_response)
+    except Exception as e:
+        logging.error(f"Failed to generate or parse slide plan: {e}")
+        return None
+
+    final_slides = []
+    for slide in slide_plan:
+        try:
+            image_url = None
+            visual_type = slide.get("visual_type")
+            visual_ref = slide.get("visual_ref")
+
+            if visual_type == "existing_chart":
+                sanitized_ref = sanitize_for_firebase_key(visual_ref)
+                image_url = chart_urls.get(sanitized_ref)
+                if not image_url:
+                    logging.warning(f"Could not find existing chart for ref: '{visual_ref}' (sanitized: '{sanitized_ref}')")
+            
+            elif visual_type == "new_image":
+                img = generate_image_with_gemini(visual_ref)
+                if img:
+                    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp_file:
+                        img_path = Path(temp_file.name)
+                        img.save(img_path, format="PNG")
+                        
+                        blob_name = f"sozo_projects/{uid}/{project_id}/slides/slide_{uuid.uuid4().hex}.png"
+                        blob = bucket.blob(blob_name)
+                        blob.upload_from_filename(str(img_path))
+                        image_url = blob.public_url
+                        logging.info(f"Uploaded new slide image to {image_url}")
+                        os.unlink(img_path)
+
+            if not image_url:
+                logging.warning(f"Visual generation failed for slide {slide.get('slide_number')}. Skipping visual for this slide.")
+
+            final_slides.append({
+                "slide_number": slide.get("slide_number"),
+                "title": slide.get("title"),
+                "content": slide.get("content"),
+                "image_url": image_url or ""
+            })
+        except Exception as slide_e:
+            logging.error(f"Failed to process slide {slide.get('slide_number')}: {slide_e}")
+            continue
+            
+    return final_slides
+
 # NEW: Keyword extraction for better Pexels searches
 def extract_keywords_for_query(text: str, llm) -> str:
     prompt = f"""