Update sozo_gen.py

sozo_gen.py

@@ -86,121 +86,23 @@ 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
+    
+def detect_dataset_domain(df: pd.DataFrame) -> str:
+    """Analyzes column names to detect the dataset's primary domain."""
+    domain_keywords = {
+        "health insurance": ["charges", "bmi", "smoker", "beneficiary"],
+        "finance": ["revenue", "profit", "cost", "budget", "expense", "stock"],
+        "marketing": ["campaign", "conversion", "click", "customer", "segment"],
+        "survey": ["satisfaction", "rating", "feedback", "opinion", "score"],
+        "food": ["nutrition", "calories", "ingredients", "restaurant"]
+    }
+    columns_lower = [col.lower() for col in df.columns]
+    for domain, keywords in domain_keywords.items():
+        if any(keyword in col for col in columns_lower for keyword in keywords):
+            logging.info(f"Dataset domain detected: {domain}")
+            return domain
+    logging.info("No specific dataset domain detected, using generic terms.")
+    return "data"
 
 # NEW: Keyword extraction for better Pexels searches
 def extract_keywords_for_query(text: str, llm) -> str:
@@ -783,4 +685,121 @@ def generate_video_from_project(df: pd.DataFrame, raw_md: str, data_context: Dic
             if os.path.exists(p): os.unlink(p)
             
         return blob.public_url
-    return None
+    return None
+
+
+# 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