Update services/queue_manager.py

services/queue_manager.py

@@ -1,18 +1,14 @@
-
 import asyncio
 import uuid
 from typing import Dict, Any, Optional
 from enum import Enum
-from core.script_gen import generate_script
-from core.story_script import generate_story
-# from core.image_gen import generate_images
-# from core.video_gen import generate_video
-# from core.music_gen import generate_music
-# from core.assemble import assemble_final_video
-
-# -------------------------------------------------------------
+
+# Import pipeline that handles all stages
+from core.pipeline import run_pipeline  # This runs story → images → video → etc.
+
+# -------------------------------
 # ENUMS AND GLOBALS
-# -------------------------------------------------------------
+# -------------------------------
 class TaskStatus(str, Enum):
     PENDING = "pending"
     RUNNING = "running"
@@ -23,10 +19,11 @@ class TaskStatus(str, Enum):
 
 tasks: Dict[str, Dict[str, Any]] = {}
 task_queue = asyncio.Queue()
+pending_confirmations: Dict[str, asyncio.Event] = {}
 
-# -------------------------------------------------------------
-# ADD NEW TASK
-# -------------------------------------------------------------
+# -------------------------------
+# ADD TASK
+# -------------------------------
 async def add_task(idea: str) -> str:
     task_id = str(uuid.uuid4())
     tasks[task_id] = {
@@ -40,56 +37,46 @@ async def add_task(idea: str) -> str:
     print(f"🧩 Task added to queue: {task_id}")
     return task_id
 
-# -------------------------------------------------------------
-# WAIT FOR SCRIPT FOR CONFIRMATION
-# -------------------------------------------------------------
+# -------------------------------
+# WAIT FOR SCRIPT
+# -------------------------------
 async def wait_for_script(task_id: str, script_results: dict):
     task = tasks.get(task_id)
     if not task:
         return
 
     task["status"] = TaskStatus.RUNNING
-    # Generate script
+
+    # === STEP 1: generate script ===
+    from core.script_gen import generate_script
     script_result = await generate_script(task["idea"])
     task["result"] = {"script": script_result}
     task["status"] = TaskStatus.WAITING_CONFIRMATION
     task["confirmation_required"] = True
 
-    # Keep script accessible for server endpoint
+    # Store script for server endpoint
     script_results[task_id] = script_result
+
+    # Create event for confirmation
+    pending_confirmations[task_id] = asyncio.Event()
     print(f"✋ Task {task_id} waiting for confirmation. Script ready.")
 
-# -------------------------------------------------------------
-# GENERATE STORY AFTER CONFIRMATION
-# -------------------------------------------------------------
-async def generate_story_after_confirm(script: str):
-    story_result = await generate_story(script)
-    return story_result
+    # Wait for user confirmation
+    await pending_confirmations[task_id].wait()
 
-# -------------------------------------------------------------
-# WORKER LOOP (Optional: future stages)
-# -------------------------------------------------------------
-async def worker():
-    while True:
-        task_id = await task_queue.get()
-        task = tasks.get(task_id)
-        if not task:
-            task_queue.task_done()
-            continue
-        try:
-            # Already handled script in wait_for_script
-            # Future stages like images/video/music can go here
-            pass
-        except Exception as e:
-            task["status"] = TaskStatus.FAILED
-            task["result"] = {"error": str(e)}
-            print(f"❌ Task {task_id} failed with error: {e}")
-        finally:
-            task_queue.task_done()
+    # Once confirmed, run full pipeline
+    task["status"] = TaskStatus.CONFIRMED
+    print(f"🎬 Task {task_id} confirmed. Running full pipeline...")
+
+    # Run pipeline (story, images, video, etc.)
+    final_result = await run_pipeline(task)
+    task["result"].update(final_result)
+    task["status"] = TaskStatus.COMPLETED
+    print(f"✅ Task {task_id} completed. All stages done.")
 
-# -------------------------------------------------------------
+# -------------------------------
 # CONFIRM TASK
-# -------------------------------------------------------------
+# -------------------------------
 async def confirm_task(task_id: str):
     task = tasks.get(task_id)
     if not task:
@@ -98,20 +85,35 @@ async def confirm_task(task_id: str):
     if task["status"] != TaskStatus.WAITING_CONFIRMATION:
         return {"error": "Task is not waiting for confirmation."}
 
-    task["status"] = TaskStatus.CONFIRMED
-    task["confirmation_required"] = False
-    print(f"👍 Task {task_id} confirmed. Ready for story generation...")
+    # Set event so the worker can continue
+    event = pending_confirmations.get(task_id)
+    if event:
+        event.set()
+
     return {"message": f"Task {task_id} confirmed."}
 
-# -------------------------------------------------------------
-# GET TASK STATUS
-# -------------------------------------------------------------
+# -------------------------------
+# GET STATUS
+# -------------------------------
 def get_task_status(task_id: str) -> Optional[Dict[str, Any]]:
     return tasks.get(task_id)
 
-# -------------------------------------------------------------
+# -------------------------------
+# WORKER
+# -------------------------------
+async def worker():
+    while True:
+        task_id = await task_queue.get()
+        task = tasks.get(task_id)
+        if not task:
+            task_queue.task_done()
+            continue
+        # Script is generated in wait_for_script, rest handled by pipeline
+        task_queue.task_done()
+
+# -------------------------------
 # START WORKER
-# -------------------------------------------------------------
+# -------------------------------
 def start_worker():
     loop = asyncio.get_event_loop()
     loop.create_task(worker())