Update main.py

main.py

@@ -20,6 +20,8 @@ import zipfile
 from fpdf import FPDF
 import tempfile
 import urllib.parse
+from stories import generateResponse
+
 
 # Initialize Flask app and CORS
 app = Flask(__name__)
@@ -247,7 +249,7 @@ def generate_story_endpoint():
             return jsonify({'error': 'Invalid or expired token'}), 401
 
         # --- Read Request Data ---
-        data = request.form.to_dict()  # For multipart/form-data fields
+        data = request.form.to_dict()  # For multipart/form-data
         input_type = data.get('input_type', 'text')  # "text", "pdf", "wiki", "bible", "youtube", "dataframe"
         prompt = data.get('prompt')  # For "text" only
         story_type = data.get('story_type', 'free_form')
@@ -256,20 +258,21 @@ def generate_story_endpoint():
         image_model = data.get('image_model', 'hf')
         audio_model = data.get('audio_model', 'deepgram')
 
-        # Validate if needed
         if input_type not in ["text", "pdf", "wiki", "bible", "youtube", "dataframe"]:
             return jsonify({'error': 'Unsupported input_type'}), 400
 
-        from stories import generate_story_from_text
-        from stories import get_pdf_text
-        from stories import get_df
-
+        from stories import (
+            generate_story_from_text,
+            get_pdf_text,
+            get_df,
+            generate_story_from_dataframe,
+            generateResponse  # <-- for chart images
+        )
         story_gen_start = time.time()
         full_story = None
 
+        # 1) Generate the full story text
         if input_type == "text":
-            # <-- CHANGE HERE
-            # We only check 'prompt' if input_type == "text"
             if not prompt:
                 return jsonify({'error': 'Prompt is required for text input'}), 400
             full_story = generate_story_from_text(prompt, story_type)
@@ -283,14 +286,14 @@ def generate_story_endpoint():
 
         elif input_type == "dataframe":
             uploaded_file = request.files.get("file")
-            ext = data.get("ext")  # e.g. "csv", "xlsx", "xls"
+            ext = data.get("ext")  # "csv", "xlsx", "xls"
             if not uploaded_file or not ext:
                 return jsonify({'error': 'File and ext are required for dataframe input'}), 400
 
             df = get_df(uploaded_file, ext)
             if df is None:
                 return jsonify({'error': f'Failed to read {ext} file'}), 400
-            from stories import generate_story_from_dataframe
+
             full_story = generate_story_from_dataframe(df, story_type)
 
         elif input_type == "wiki":
@@ -321,7 +324,7 @@ def generate_story_endpoint():
         if not full_story:
             return jsonify({'error': 'Story generation failed'}), 500
 
-        # 2) Split the story into 5 sections
+        # 2) Split into 5 sections
         sections_raw = [s.strip() for s in full_story.split("[break]") if s.strip()]
         if len(sections_raw) < 5:
             sections_raw += ["(Placeholder section)"] * (5 - len(sections_raw))
@@ -335,33 +338,55 @@ def generate_story_endpoint():
         from image_gen import generate_image_with_retry
         from audio_gen import generate_audio
 
+        # If input_type is "dataframe", we have a df for chart generation
+        df = None
+        if input_type == "dataframe":
+            uploaded_file = request.files.get("file")
+            ext = data.get("ext")
+            df = get_df(uploaded_file, ext)  # re-use the same df
+
         # 3) Process each section
         for section_text in sections_raw:
             # Extract an image prompt between angle brackets
             img_prompt_match = re.search(r"<(.*?)>", section_text)
             img_prompt = img_prompt_match.group(1).strip() if img_prompt_match else section_text[:100]
 
-            # Generate image
             image_start = time.time()
-            image_obj, _ = generate_image_with_retry(img_prompt, style, model=image_model)
+            image_obj = None
+
+            # If we are dealing with "dataframe", attempt chart generation first
+            if input_type == "dataframe" and df is not None:
+                try:
+                    chart_str = generateResponse(img_prompt, df)  # returns a Python string or None
+                    if chart_str and chart_str.startswith("data:image/png;base64,"):
+                        # decode base64 -> PIL Image
+                        base64_data = chart_str.split(",", 1)[1]
+                        chart_bytes = base64.b64decode(base64_data)
+                        image_obj = Image.open(io.BytesIO(chart_bytes))
+                except Exception as e:
+                    print("DataFrame chart generation error:", e)
+
+            # Fallback to generate_image_with_retry
+            if not image_obj:
+                image_obj, _ = generate_image_with_retry(img_prompt, style, model=image_model)
+
             image_end = time.time()
             image_generation_times.append(image_end - image_start)
 
-            # Save image locally -> upload -> get URL
+            # Save & upload
             image_filename = f"/tmp/{uuid.uuid4().hex}.jpg"
             image_obj.save(image_filename, format="JPEG")
             image_blob_name = f"stories/{uid}/{uuid.uuid4().hex}.jpg"
             image_url = upload_to_storage(image_filename, image_blob_name)
             os.remove(image_filename)
 
-            # Generate audio from section text WITHOUT <image> description
-            audio_text = re.sub(r"<.*?>", "", section_text)  # remove anything in angle brackets
+            # Generate audio without <image> description
+            audio_text = re.sub(r"<.*?>", "", section_text)
             audio_start = time.time()
             audio_file_path = generate_audio(audio_text, voice_model, audio_model=audio_model)
             audio_end = time.time()
             audio_generation_times.append(audio_end - audio_start)
 
-            # Upload audio
             audio_blob_name = f"stories/{uid}/{uuid.uuid4().hex}.mp3"
             audio_url = upload_to_storage(audio_file_path, audio_blob_name)
             os.remove(audio_file_path)
@@ -372,7 +397,7 @@ def generate_story_endpoint():
                 "audio_url": audio_url
             })
 
-        # 4) Store the story record in Firebase Realtime Database
+        # 4) Store the story
         story_id = str(uuid.uuid4())
         story_ref = db.reference(f"stories/{story_id}")
         story_record = {