Update sozo_gen.py

sozo_gen.py

@@ -502,7 +502,7 @@ def get_augmented_context(df: pd.DataFrame, user_ctx: str) -> Dict:
 def generate_report_draft(buf, name: str, ctx: str, uid: str, project_id: str, bucket):
     logging.info(f"Generating guided storyteller report draft for project {project_id}")
     df = load_dataframe_safely(buf, name)
-    llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash", google_api_key=API_KEY, temperature=0.25)
+    llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash", google_api_key=API_KEY, temperature=0.3)
     
     data_context_str, context_for_charts = "", {}
     try:
@@ -522,20 +522,27 @@ def generate_report_draft(buf, name: str, ctx: str, uid: str, project_id: str, b
 
     md = ""
     try:
-        # --- Pass 1: The Visualization Strategist ---
+        # --- Pass 1: The "Visualization Strategist" ---
         strategist_prompt = f"""
-        You are a data visualization expert. Your task is to create a palette of unique and impactful charts for a data storyteller.
+        You are a data visualization expert. Your task is to create a diverse palette of unique and impactful charts for a data storyteller.
         Based on the provided data context, identify the 4-5 most distinct and insightful stories that can be visualized.
-        For each one, provide only the chart description tag.
         
         **Data Context:**
         {data_context_str}
 
+        **Your Goal:**
+        Your primary goal is to select a **diverse palette of chart types**. A high-quality response will use a mix of different charts from the available list to create a visually engaging and comprehensive report. **Do not use the same chart type more than twice.**
+
+        **Strategic Hints:**
+        - Consider a `histogram` to show the distribution of a key variable (like age or bmi).
+        - Consider a `pie chart` for a clear part-to-whole relationship (e.g., smoker vs. non-smoker proportions).
+        - Consider a `heatmap` if the dataset has multiple numeric columns and you believe the overall pattern of their correlations is a key insight in itself.
+
         **Output Format:**
         Return ONLY a valid JSON array of strings. Each string must be a unique chart description tag.
         
         Example:
-        ["bar | Average Charges by Smoker Status", "scatter | Charges vs. BMI", "bar | Average Charges by Region"]
+        ["bar | Average Charges by Smoker Status", "scatter | Charges vs. BMI", "hist | Distribution of Beneficiary Ages", "pie | Regional Proportions"]
         """
         logging.info("Executing Visualization Strategist Pass...")
         strategist_response = llm.invoke(strategist_prompt).content.strip()
@@ -544,13 +551,18 @@ def generate_report_draft(buf, name: str, ctx: str, uid: str, project_id: str, b
         chart_palette = json.loads(strategist_response)
         logging.info(f"Strategist Pass successful. Palette has {len(chart_palette)} unique charts.")
 
-        # --- Pass 2: The Master Storyteller ---
+        # --- Pass 2: The "Master Storyteller" ---
         storyteller_prompt = f"""
-        You are an elite data storyteller. Your mission is to write a comprehensive, flowing narrative that analyzes the entire dataset provided. Weave a story that connects different findings, explores nuances, and provides a holistic view.
+        You are an elite data storyteller and business intelligence expert. Your mission is to write a comprehensive, flowing narrative that analyzes the entire dataset provided. Your goal is to create a captivating story that **drives action**.
 
         **Data Context:**
         {data_context_str}
 
+        **Narrative Construction Guidelines:**
+        1.  **Use Compelling Headers:** Structure your report with multiple sections using Markdown headings (`##` or `###`). Do not write one long block of text. Create curiosity with your headers (e.g., 'The Smoking Premium: A Costly Habit', 'Geographic Hotspots: Where Charges Are Highest').
+        2.  **Weave a Story:** Don't just describe the charts one by one. Connect the findings together. For example, how does 'age' relate to 'smoker status' and how do they both impact 'charges'?
+        3.  **Drive to Action:** Conclude your report with a dedicated section titled `## Actionable Recommendations`. Based on your analysis, provide specific, data-driven suggestions that a business leader could implement.
+
         **Your Toolbox (Most Important):**
         To support your story with visuals, you have been provided with a pre-approved 'palette' of unique charts. As you write your narrative, you **must** integrate each of these chart tags, one time, at the most logical point in the story.
         - You **must** use every chart tag from the provided palette exactly once.