Update sozo_gen.py

sozo_gen.py

@@ -175,41 +175,42 @@ class ChartGenerator:
 
 # --- Main Business Logic Functions ---
 
-def generate_report_draft(df: pd.DataFrame, user_context: str) -> Dict[str, Any]:
-    """
-    Generates a markdown report draft and a list of chart specifications.
-    This function does NOT interact with storage.
-    """
-    logging.info("Generating report draft and chart specifications.")
-    llm = ChatGoogleGenerativeAI(model="gemini-1.5-flash", google_api_key=API_KEY, temperature=0.1)
-    ctx_dict = {"shape": df.shape, "columns": list(df.columns), "user_ctx": user_context}
-    enhanced_ctx = enhance_data_context(df, ctx_dict)
+# In sozo_gen.py
 
+def generate_report_draft(buf, name: str, ctx: str, uid: str, project_id: str, bucket):
+    logging.info(f"Generating report draft for project {project_id}")
+    df = load_dataframe_safely(buf, name)
+    llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", google_api_key=API_KEY, temperature=0.1)
+    ctx_dict = {"shape": df.shape, "columns": list(df.columns), "user_ctx": ctx}
+    enhanced_ctx = enhance_data_context(df, ctx_dict)
     report_prompt = f"""
-    You are a senior data analyst. Analyze the provided dataset context and write a comprehensive executive-level Markdown report.
+    You are a senior data analyst and business intelligence expert. Analyze the provided dataset and write a comprehensive executive-level Markdown report.
     **Dataset Analysis Context:** {json.dumps(enhanced_ctx, indent=2)}
     **Instructions:**
     1. **Executive Summary**: Start with a high-level summary of key findings.
-    2. **Key Insights**: Provide 3-5 key insights. For each insight that can be visualized, insert a chart tag on its own line.
-    3. **Visual Support**: Use chart tags like: <generate_chart: "A bar chart showing total sales by region">.
+    2. **Key Insights**: Provide 3-5 key insights, each with its own chart tag.
+    3. **Visual Support**: Insert chart tags like: `<generate_chart: "chart_type | specific description">`.
        Valid chart types: bar, pie, line, scatter, hist.
     Generate insights that would be valuable to C-level executives.
     """
     md = llm.invoke(report_prompt).content
     chart_descs = extract_chart_tags(md)[:MAX_CHARTS]
-    chart_specs = []
-
+    chart_urls = {}
     chart_generator = ChartGenerator(llm, df)
     for desc in chart_descs:
-        try:
-            spec = chart_generator.generate_chart_spec(desc)
-            chart_specs.append(spec)
-        except Exception as e:
-            logging.warning(f"Could not generate spec for chart '{desc}': {e}")
-            # Continue without the failed chart
-
-    return {"raw_md": md, "chart_specs": chart_specs}
-
+        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp_file:
+            img_path = Path(temp_file.name)
+            try:
+                chart_spec = chart_generator.generate_chart_spec(desc)
+                if execute_chart_spec(chart_spec, df, img_path):
+                    blob_name = f"sozo_projects/{uid}/{project_id}/charts/{uuid.uuid4().hex}.png"
+                    blob = bucket.blob(blob_name)
+                    blob.upload_from_filename(str(img_path))
+                    chart_urls[desc] = blob.public_url
+                    logging.info(f"Uploaded chart '{desc}' to {blob.public_url}")
+            finally:
+                os.unlink(img_path)
+    return {"raw_md": md, "chartUrls": chart_urls}
 
 def generate_single_chart(df: pd.DataFrame, description: str) -> Dict[str, Any]:
     """Generates a single chart specification dictionary."""