Update main.py

main.py

@@ -247,67 +247,86 @@ def generate_story_endpoint():
             return jsonify({'error': 'Invalid or expired token'}), 401
 
         # --- Read Request Data ---
-        data = request.get_json()
-        input_type = data.get('input_type', 'text')  # "text", "wiki", "bible", "youtube", or "dataframe"
-        prompt = data.get('prompt')
+        # If the user is uploading a file (PDF or CSV/Excel), we can read from request.files
+        # If the user is sending JSON only, we read request.get_json()
+        data = request.form.to_dict()  # For multipart/form-data fields
+        input_type = data.get('input_type', 'text')  # "text", "pdf", "wiki", "bible", "youtube", "dataframe"
+        prompt = data.get('prompt')  # For "text" or fallback
         story_type = data.get('story_type', 'free_form')
         style = data.get('style', 'whimsical')
         voice_model = data.get('voice_model', 'aura-asteria-en')
         image_model = data.get('image_model', 'hf')
         audio_model = data.get('audio_model', 'deepgram')
 
-        if not prompt:
-            return jsonify({'error': 'Prompt is required'}), 400
+        # Validate if needed
+        if input_type not in ["text", "pdf", "wiki", "bible", "youtube", "dataframe"]:
+            return jsonify({'error': 'Unsupported input_type'}), 400
+
+        # 1) Generate the full story text
+        from stories import generate_story_from_text
+        from stories import get_pdf_text
+        from stories import get_df
 
-      
-        # --- Select the Appropriate Story Generation Function ---
         story_gen_start = time.time()
         full_story = None
 
         if input_type == "text":
-            from stories import generate_story_from_text
-            full_story = generate_story_from_text(prompt, story_type)
-        if input_type == "pdf":
-            from stories import generate_story_from_text
-            from stories import get_pdf_text
-            prompt = get_pdf_text(pdf)
+            if not prompt:
+                return jsonify({'error': 'Prompt is required for text input'}), 400
             full_story = generate_story_from_text(prompt, story_type)
+
+        elif input_type == "pdf":
+            # Expecting a file in request.files["file"]
+            uploaded_file = request.files.get("file")
+            if not uploaded_file:
+                return jsonify({'error': 'No PDF file uploaded'}), 400
+
+            # Convert PDF to text
+            pdf_text = get_pdf_text(uploaded_file)
+            full_story = generate_story_from_text(pdf_text, story_type)
+
+        elif input_type == "dataframe":
+            # Expecting a file in request.files["file"] and an "ext" field (csv, xlsx, xls)
+            uploaded_file = request.files.get("file")
+            ext = data.get("ext")  # e.g. "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":
             wiki_url = data.get("wiki_url")
             if not wiki_url:
                 return jsonify({'error': 'wiki_url is required for input_type "wiki"'}), 400
             from stories import generate_story_from_wiki
             full_story = generate_story_from_wiki(wiki_url, story_type)
+
         elif input_type == "bible":
             bible_reference = data.get("bible_reference")
             if not bible_reference:
                 return jsonify({'error': 'bible_reference is required for input_type "bible"'}), 400
             from stories import generate_story_from_bible
             full_story = generate_story_from_bible(bible_reference, story_type)
+
         elif input_type == "youtube":
             youtube_url = data.get("youtube_url")
             if not youtube_url:
                 return jsonify({'error': 'youtube_url is required for input_type "youtube"'}), 400
             from stories import generate_story_from_youtube
             full_story = generate_story_from_youtube(youtube_url, story_type)
-        elif input_type == "dataframe":
-            # Expecting dataframe data as JSON (list of dicts)
-            df_data = data.get("data")
-            if not df_data:
-                return jsonify({'error': 'Data for dataframe input_type is required'}), 400
-            df = pd.DataFrame(df_data)
-            from stories import generate_story_from_dataframe
-            full_story = generate_story_from_dataframe(df, story_type)
-        else:
-            return jsonify({'error': 'Unsupported input_type'}), 400
 
+        # Measure generation time
         story_gen_end = time.time()
         story_generation_time = story_gen_end - story_gen_start
 
         if not full_story:
             return jsonify({'error': 'Story generation failed'}), 500
 
-        # --- Split the Story into 5 Sections ---
+        # 2) Split the story 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))
@@ -318,48 +337,48 @@ def generate_story_endpoint():
         image_generation_times = []
         audio_generation_times = []
 
-        # Import generation functions from your modules
-        from image_gen import generate_image_with_retry  # image generation function
-        from audio_gen import generate_audio              # audio generation function
+        from image_gen import generate_image_with_retry
+        from audio_gen import generate_audio
 
-        # Process each section
-        for section in sections_raw:
-            # --- Image Generation ---
-            # Extract an image prompt between angle brackets; otherwise, fallback to the first 100 characters.
-            img_prompt_match = re.search(r"<(.*?)>", section)
-            img_prompt = img_prompt_match.group(1).strip() if img_prompt_match else section[:100]
+        # 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_end = time.time()
             image_generation_times.append(image_end - image_start)
 
-            # Save image locally and upload it.
+            # Save image locally -> upload -> get URL
             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)
 
-            # --- Audio Generation ---
+            # Generate audio from section text WITHOUT <image> description
+            # e.g. remove <...> from text
+            audio_text = re.sub(r"<.*?>", "", section_text)  # remove anything in angle brackets
             audio_start = time.time()
-            audio_file_path = generate_audio(section, voice_model, audio_model=audio_model)
+            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)
 
             sections.append({
-                "section_text": section,
+                "section_text": section_text,
                 "image_url": image_url,
                 "audio_url": audio_url
             })
 
-            # Clean up temporary files
-            os.remove(image_filename)
-            os.remove(audio_file_path)
-
-        # --- Store the Story Record in Firebase Realtime Database ---
+        # 4) Store the story record in Firebase Realtime Database
         story_id = str(uuid.uuid4())
         story_ref = db.reference(f"stories/{story_id}")
         story_record = {
@@ -376,7 +395,8 @@ def generate_story_endpoint():
             "story_type": story_type
         }
         story_ref.set(story_record)
-          # --- Subtract 5 Credits from the User ---
+
+        # Subtract 5 Credits
         user_ref = db.reference(f"users/{uid}")
         user_data = user_ref.get() or {}
         current_credits = user_data.get("credits", 0)