Update sozo_gen.py

sozo_gen.py

@@ -106,9 +106,10 @@ class ChartGenerator:
         self.enhanced_ctx = enhance_data_context(df, {"columns": list(df.columns), "shape": df.shape, "dtypes": {col: str(dtype) for col, dtype in df.dtypes.items()}})
 
     def generate_chart_spec(self, description: str) -> ChartSpecification:
+        safe_ctx = json_serializable(self.enhanced_ctx)
         spec_prompt = f"""
         You are a data visualization expert. Based on the dataset and chart description, generate a precise chart specification.
-        **Dataset Info:** {json.dumps(self.enhanced_ctx, indent=2)}
+        **Dataset Info:** {json.dumps(safe_ctx, indent=2)}
         **Chart Request:** {description}
         **Return a JSON specification with these exact fields:**
         {{
@@ -224,7 +225,7 @@ def animate_image_fade(img: np.ndarray, dur: float, out: Path, fps: int = 24) ->
 
 def safe_chart(desc: str, df: pd.DataFrame, dur: float, out: Path) -> str:
     try:
-        llm = ChatGoogleGenerativeAI(model="gemini-1.5-flash-latest", google_api_key=API_KEY, temperature=0.1)
+        llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", google_api_key=API_KEY, temperature=0.1)
         chart_generator = ChartGenerator(llm, df)
         chart_spec = chart_generator.generate_chart_spec(desc)
         return animate_chart(chart_spec, df, dur, out)
@@ -232,7 +233,7 @@ def safe_chart(desc: str, df: pd.DataFrame, dur: float, out: Path) -> str:
         logging.error(f"Chart animation failed for '{desc}': {e}. Falling back to static image.")
         with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp_png_file:
             temp_png = Path(temp_png_file.name)
-        llm = ChatGoogleGenerativeAI(model="gemini-1.5-flash-latest", google_api_key=API_KEY, temperature=0.1)
+        llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", google_api_key=API_KEY, temperature=0.1)
         chart_generator = ChartGenerator(llm, df)
         chart_spec = chart_generator.generate_chart_spec(desc)
         if execute_chart_spec(chart_spec, df, temp_png):