rairo commited on
Commit
ea70ca1
·
verified ·
1 Parent(s): 848688d

Update sozo_gen.py

Browse files
Files changed (1) hide show
  1. sozo_gen.py +23 -3
sozo_gen.py CHANGED
@@ -259,6 +259,15 @@ def concat_media(file_paths: List[str], output_path: Path):
259
 
260
  # --- Main Business Logic Functions for Flask ---
261
 
 
 
 
 
 
 
 
 
 
262
  def generate_report_draft(buf, name: str, ctx: str, uid: str, project_id: str, bucket):
263
  logging.info(f"Generating report draft for project {project_id}")
264
  df = load_dataframe_safely(buf, name)
@@ -279,20 +288,31 @@ def generate_report_draft(buf, name: str, ctx: str, uid: str, project_id: str, b
279
  chart_descs = extract_chart_tags(md)[:MAX_CHARTS]
280
  chart_urls = {}
281
  chart_generator = ChartGenerator(llm, df)
 
282
  for desc in chart_descs:
 
 
 
 
 
 
 
283
  with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp_file:
284
  img_path = Path(temp_file.name)
285
  try:
286
- chart_spec = chart_generator.generate_chart_spec(desc)
287
  if execute_chart_spec(chart_spec, df, img_path):
288
  blob_name = f"sozo_projects/{uid}/{project_id}/charts/{uuid.uuid4().hex}.png"
289
  blob = bucket.blob(blob_name)
290
  blob.upload_from_filename(str(img_path))
291
- chart_urls[desc] = blob.public_url
292
- logging.info(f"Uploaded chart '{desc}' to {blob.public_url}")
 
 
293
  finally:
294
  if os.path.exists(img_path):
295
  os.unlink(img_path)
 
296
  return {"raw_md": md, "chartUrls": chart_urls}
297
 
298
  def generate_single_chart(df: pd.DataFrame, description: str, uid: str, project_id: str, bucket):
 
259
 
260
  # --- Main Business Logic Functions for Flask ---
261
 
262
+ # ADD THIS NEW HELPER FUNCTION SOMEWHERE NEAR THE TOP OF THE FILE
263
+ def sanitize_for_firebase_key(text: str) -> str:
264
+ """Replaces Firebase-forbidden characters in a string with underscores."""
265
+ forbidden_chars = ['.', '$', '#', '[', ']', '/']
266
+ for char in forbidden_chars:
267
+ text = text.replace(char, '_')
268
+ return text
269
+
270
+ # REPLACE THE OLD generate_report_draft WITH THIS CORRECTED VERSION
271
  def generate_report_draft(buf, name: str, ctx: str, uid: str, project_id: str, bucket):
272
  logging.info(f"Generating report draft for project {project_id}")
273
  df = load_dataframe_safely(buf, name)
 
288
  chart_descs = extract_chart_tags(md)[:MAX_CHARTS]
289
  chart_urls = {}
290
  chart_generator = ChartGenerator(llm, df)
291
+
292
  for desc in chart_descs:
293
+ # Create a safe key for Firebase
294
+ safe_desc = sanitize_for_firebase_key(desc)
295
+
296
+ # Replace the original description in the markdown with the safe one
297
+ md = md.replace(f'<generate_chart: "{desc}">', f'<generate_chart: "{safe_desc}">')
298
+ md = md.replace(f'<generate_chart: {desc}>', f'<generate_chart: "{safe_desc}">') # Handle no quotes case
299
+
300
  with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp_file:
301
  img_path = Path(temp_file.name)
302
  try:
303
+ chart_spec = chart_generator.generate_chart_spec(desc) # Still generate spec from original desc
304
  if execute_chart_spec(chart_spec, df, img_path):
305
  blob_name = f"sozo_projects/{uid}/{project_id}/charts/{uuid.uuid4().hex}.png"
306
  blob = bucket.blob(blob_name)
307
  blob.upload_from_filename(str(img_path))
308
+
309
+ # Use the safe key in the dictionary
310
+ chart_urls[safe_desc] = blob.public_url
311
+ logging.info(f"Uploaded chart '{desc}' to {blob.public_url} with safe key '{safe_desc}'")
312
  finally:
313
  if os.path.exists(img_path):
314
  os.unlink(img_path)
315
+
316
  return {"raw_md": md, "chartUrls": chart_urls}
317
 
318
  def generate_single_chart(df: pd.DataFrame, description: str, uid: str, project_id: str, bucket):