Update sozo_gen.py

sozo_gen.py

@@ -86,6 +86,15 @@ def clean_narration(txt: str) -> str:
 
 def placeholder_img() -> Image.Image: return Image.new("RGB", (WIDTH, HEIGHT), (230, 230, 230))
 
+def _sanitize_for_json(data):
+    """Recursively sanitizes a dict/list for JSON compliance."""
+    if isinstance(data, dict):
+        return {k: _sanitize_for_json(v) for k, v in data.items()}
+    if isinstance(data, list):
+        return [_sanitize_for_json(i) for i in data]
+    if isinstance(data, float) and (math.isnan(data) or math.isinf(data)):
+        return None
+    return data
     
 def detect_dataset_domain(df: pd.DataFrame) -> str:
     """Analyzes column names to detect the dataset's primary domain."""
@@ -415,20 +424,33 @@ def generate_visualization_strategy(intelligence: Dict) -> str:
     return strategy
 
 def get_augmented_context(df: pd.DataFrame, user_ctx: str) -> Dict:
+    """Creates a detailed, JSON-safe summary of the dataframe for the AI."""
     numeric_cols = df.select_dtypes(include=['number']).columns.tolist()
     categorical_cols = df.select_dtypes(exclude=['number']).columns.tolist()
+    
     context = {
         "user_context": user_ctx,
         "dataset_shape": {"rows": df.shape[0], "columns": df.shape[1]},
         "schema": {"numeric_columns": numeric_cols, "categorical_columns": categorical_cols},
         "data_previews": {}
     }
+    
     for col in categorical_cols[:5]:
         unique_vals = df[col].unique()
-        context["data_previews"][col] = {"count": len(unique_vals), "values": unique_vals[:5].tolist()}
+        context["data_previews"][col] = {
+            "count": len(unique_vals),
+            "values": unique_vals[:5].tolist()
+        }
+        
     for col in numeric_cols[:5]:
-        context["data_previews"][col] = {"mean": df[col].mean(), "min": df[col].min(), "max": df[col].max()}
-    return json.loads(json.dumps(context, default=str))
+        context["data_previews"][col] = {
+            "mean": df[col].mean(),
+            "min": df[col].min(),
+            "max": df[col].max()
+        }
+    
+    # Sanitize the entire structure before returning
+    return _sanitize_for_json(json.loads(json.dumps(context, default=str)))
 
 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}")