Update app.py

app.py

@@ -4,6 +4,7 @@ import httpx
 import torch
 import logging
 import re
+import json
 from typing import Dict, Optional, List, Union
 from fastapi import FastAPI, Request, BackgroundTasks, HTTPException, Depends
 from fastapi.responses import JSONResponse
@@ -84,13 +85,75 @@ async def lifespan(app: FastAPI):
 
 app = FastAPI(lifespan=lifespan)
 
-def extract_topic(topic_input: Union[str, List[str]]) -> str:
-    """Extract topic from string or array input"""
-    if isinstance(topic_input, list):
-        if topic_input:
-            return str(topic_input[0])
-        return "No topic provided"
-    return str(topic_input)
+def extract_topics(topic_input: Union[str, List[str]]) -> List[str]:
+    """Extract and validate topics from input"""
+    if isinstance(topic_input, str):
+        try:
+            # Try to parse as JSON if it's a string
+            parsed = json.loads(topic_input)
+            if isinstance(parsed, list):
+                return [str(topic).strip() for topic in parsed if str(topic).strip()]
+            return [str(parsed).strip()]
+        except json.JSONDecodeError:
+            # If not JSON, treat as comma-separated string
+            if "," in topic_input:
+                return [topic.strip() for topic in topic_input.split(",") if topic.strip()]
+            return [topic_input.strip()]
+    elif isinstance(topic_input, list):
+        return [str(topic).strip() for topic in topic_input if str(topic).strip()]
+    
+    return []
+
+def generate_topic_from_trends(trending_topics: List[str]) -> str:
+    """Generate a viral topic based on trending topics"""
+    if not generator.loaded:
+        if not generator.load_model():
+            raise Exception(f"Model failed to load: {generator.load_error}")
+    
+    logger.info(f"🧠 Generating viral topic from trends: {trending_topics}")
+    
+    prompt = (
+        f"Based on these 5 trending topics: {', '.join(trending_topics)}\n\n"
+        "Create ONE highly engaging, viral topic for a YouTube/TikTok short video that:\n"
+        "1. Combines elements from these trends in a creative way\n"
+        "2. Has high viral potential (emotional, surprising, or controversial)\n"
+        "3. Is suitable for a 60-second video format\n"
+        "4. Appeals to a broad audience\n"
+        "5. Is specific enough to be interesting but broad enough to allow creative interpretation\n\n"
+        "Respond with ONLY the topic (no explanations, no bullet points, no numbering).\n"
+        "The topic should be 5-10 words maximum.\n\n"
+        "VIRAL TOPIC:"
+    )
+    
+    inputs = generator.tokenizer(
+        prompt, 
+        return_tensors="pt", 
+        truncation=True,
+        max_length=512
+    )
+    
+    inputs = {k: v.to(DEVICE) for k, v in inputs.items()}
+    
+    with torch.no_grad():
+        outputs = generator.model.generate(
+            **inputs,
+            max_new_tokens=100,
+            do_sample=True,
+            top_p=0.9,
+            temperature=0.8,
+            pad_token_id=generator.tokenizer.eos_token_id,
+            repetition_penalty=1.1
+        )
+    
+    generated_text = generator.tokenizer.decode(outputs[0], skip_special_tokens=True)
+    topic = generated_text.replace(prompt, "").strip()
+    
+    # Clean up the topic
+    topic = re.split(r'[\n\.]', topic)[0].strip()
+    topic = re.sub(r'^["\'](.*)["\']$', r'\1', topic)  # Remove surrounding quotes
+    
+    logger.info(f"🎯 Generated topic: '{topic}'")
+    return topic
 
 def format_script(script: str) -> str:
     """Clean and format the generated script"""
@@ -200,19 +263,28 @@ def generate_script(topic: str) -> str:
         logger.error(f"❌ Script generation failed: {str(e)}")
         raise
 
-async def process_job(job_id: str, topic_input: Union[str, List[str]], callback_url: str = None):
+async def process_job(job_id: str, topics_input: Union[str, List[str]], callback_url: str = None):
     """Background task to process job"""
     try:
-        topic = extract_topic(topic_input)
-        logger.info(f"🎯 Processing: '{topic}'")
+        # Extract and validate topics
+        topics = extract_topics(topics_input)
+        if len(topics) < 3:
+            raise HTTPException(status_code=400, detail="At least 3 topics are required")
+            
+        logger.info(f"🎯 Processing {len(topics)} topics: {topics}")
         
-        script = generate_script(topic)
+        # Step 1: Generate a viral topic from the trends
+        generated_topic = generate_topic_from_trends(topics)
+        
+        # Step 2: Generate script based on the created topic
+        script = generate_script(generated_topic)
         
         # Store job results
         jobs[job_id] = {
             "status": "complete",
             "result": script,
-            "topic": topic,
+            "original_topics": topics,
+            "generated_topic": generated_topic,
             "script_length": len(script),
             "formatted": True
         }
@@ -227,7 +299,8 @@ async def process_job(job_id: str, topic_input: Union[str, List[str]], callback_
                         "job_id": job_id,
                         "status": "complete",
                         "result": script,
-                        "topic": topic,
+                        "original_topics": topics,
+                        "generated_topic": generated_topic,
                         "script_length": len(script),
                         "formatted": True
                     }
@@ -251,7 +324,7 @@ async def process_job(job_id: str, topic_input: Union[str, List[str]], callback_
         jobs[job_id] = {
             "status": "failed",
             "error": error_msg,
-            "topic": extract_topic(topic_input) if topic_input else "unknown"
+            "topics": extract_topics(topics_input) if topics_input else []
         }
         
         # Send failure webhook if callback URL exists
@@ -264,7 +337,7 @@ async def process_job(job_id: str, topic_input: Union[str, List[str]], callback_
                             "job_id": job_id,
                             "status": "failed",
                             "error": error_msg,
-                            "topic": extract_topic(topic_input) if topic_input else "unknown"
+                            "topics": extract_topics(topics_input) if topics_input else []
                         }
                     )
             except Exception:
@@ -282,36 +355,39 @@ async def submit_job(
         job_id = str(uuid.uuid4())
         
         # Validate input
-        if not data.get("topic"):
-            raise HTTPException(status_code=400, detail="Topic is required")
+        if not data.get("topics"):
+            raise HTTPException(status_code=400, detail="Topics are required")
             
         callback_url = data.get("callback_url")
-        topic_input = data["topic"]
-        topic = extract_topic(topic_input)
+        topics_input = data["topics"]
+        topics = extract_topics(topics_input)
         
-        logger.info(f"📥 Received job {job_id}: '{topic}'")
+        if len(topics) < 3:
+            raise HTTPException(status_code=400, detail="At least 3 topics are required")
+            
+        logger.info(f"📥 Received job {job_id} with {len(topics)} topics: {topics}")
         
         # Store initial job data
         jobs[job_id] = {
             "status": "processing",
             "callback_url": callback_url,
-            "topic": topic
+            "topics": topics
         }
         
         # Process job in background
         background_tasks.add_task(
             process_job,
             job_id,
-            topic_input,
+            topics_input,
             callback_url
         )
         
         return JSONResponse({
             "job_id": job_id,
             "status": "queued",
-            "topic": topic,
-            "estimated_time": "70-90 seconds",
-            "message": "Script generation started"
+            "topics": topics,
+            "estimated_time": "90-120 seconds",
+            "message": "Topic generation and script creation started"
         })
         
     except Exception as e:
@@ -348,14 +424,26 @@ async def test_generation(auth: bool = Depends(verify_api_key)):
             if not generator.load_model():
                 return {"status": "error", "error": "Model failed to load"}
         
-        test_topic = "the future of artificial intelligence in healthcare"
-        logger.info(f"🧪 Testing generation with: {test_topic}")
+        test_topics = [
+            "Artificial Intelligence", 
+            "Sustainable Energy", 
+            "Virtual Reality", 
+            "Space Exploration", 
+            "Biotechnology"
+        ]
+        
+        logger.info(f"🧪 Testing topic generation with: {test_topics}")
+        
+        # Test topic generation
+        generated_topic = generate_topic_from_trends(test_topics)
         
-        script = generate_script(test_topic)
+        # Test script generation
+        script = generate_script(generated_topic)
         
         return {
             "status": "success", 
-            "topic": test_topic, 
+            "test_topics": test_topics,
+            "generated_topic": generated_topic,
             "script_length": len(script),
             "script_preview": script[:300] + "..." if len(script) > 300 else script,
             "estimated_duration": "60 seconds",
@@ -370,10 +458,11 @@ async def test_generation(auth: bool = Depends(verify_api_key)):
 async def root():
     """Root endpoint"""
     return {
-        "message": "Video Script Generator API",
-        "version": "1.0",
+        "message": "Enhanced Video Script Generator API",
+        "version": "2.0",
+        "features": "Generates viral topics from trends, then creates scripts",
         "endpoints": {
-            "submit_job": "POST /api/submit",
+            "submit_job": "POST /api/submit (with 'topics' array)",
             "check_status": "GET /api/status/{job_id}",
             "health": "GET /health",
             "test": "GET /test/generation"