Update app.py

app.py

@@ -3,6 +3,7 @@ import uuid
 import httpx
 import torch
 import logging
+import socket
 from typing import Dict, Optional
 from fastapi import FastAPI, Request, BackgroundTasks, HTTPException
 from fastapi.responses import JSONResponse
@@ -10,6 +11,7 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 import uvicorn
 from contextlib import asynccontextmanager
 from pydantic import BaseModel
+from tenacity import retry, stop_after_attempt, wait_exponential
 
 # Configuration
 MODEL_ID = "google/gemma-1.1-2b-it"
@@ -56,17 +58,18 @@ class ScriptGenerator:
             ).to(DEVICE)
             self.loaded = True
             logger.info("Model loaded successfully")
-            self.send_telegram_notification("🤖 Model loaded and ready to generate scripts!")
+            self.send_telegram_notification("🤖 Model loaded and ready!")
         except Exception as e:
             error_msg = f"Model loading failed: {str(e)}"
             logger.error(error_msg)
             self.send_telegram_notification(f"❌ {error_msg}")
             raise
 
+    @retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10))
     async def send_telegram_notification(self, message: str):
-        """Send notification to Telegram"""
+        """Send notification to Telegram with retry logic"""
         if not TELEGRAM_BOT_TOKEN or not TELEGRAM_CHAT_ID:
-            logger.warning("Telegram credentials not configured")
+            logger.warning("Telegram not configured")
             return
             
         try:
@@ -77,11 +80,17 @@ class ScriptGenerator:
                 "parse_mode": "Markdown"
             }
             
-            async with httpx.AsyncClient() as client:
-                await client.post(url, json=payload, timeout=10.0)
+            transport = httpx.AsyncHTTPTransport(retries=3)
+            timeout = httpx.Timeout(15.0, connect=10.0)
+            
+            async with httpx.AsyncClient(transport=transport, timeout=timeout) as client:
+                response = await client.post(url, json=payload)
+                response.raise_for_status()
+                logger.info(f"Telegram notification sent")
                 
         except Exception as e:
             logger.error(f"Telegram notification failed: {str(e)}")
+            raise
 
 # Lifespan handler
 @asynccontextmanager
@@ -89,11 +98,18 @@ async def lifespan(app: FastAPI):
     generator = ScriptGenerator()
     generator.load_model()
     yield
-    # Cleanup if needed
 
 app = FastAPI(lifespan=lifespan)
 generator = ScriptGenerator()
 
+def check_dns():
+    """Check DNS resolution"""
+    try:
+        socket.gethostbyname('api.telegram.org')
+        return True
+    except socket.gaierror:
+        return False
+
 def generate_script(topic: str) -> str:
     """Generate script with error handling"""
     try:
@@ -101,11 +117,11 @@ def generate_script(topic: str) -> str:
             generator.load_model()
             
         prompt = (
-            f"Create a 1-minute engaging video script about: {topic[:80]}\n\n"
+            f"Create a 1-minute video script about: {topic[:80]}\n\n"
             "Structure:\n"
-            "1) Hook (5-10 seconds - grab attention)\n"
-            "2) Main Content (40-45 seconds - valuable information)\n"
-            "3) CTA (5-10 seconds - clear call-to-action)\n\n"
+            "1) Hook (5-10 seconds)\n"
+            "2) Main Content (40-45 seconds)\n"
+            "3) CTA (5-10 seconds)\n\n"
             "Script:"
         )
         
@@ -130,132 +146,78 @@ async def process_job(job_id: str, topic: str, callback_url: str = None):
     """Background task to process job"""
     try:
         await generator.send_telegram_notification(
-            f"⏳ *Job Started*\n"
-            f"ID: `{job_id}`\n"
-            f"Topic: _{topic}_"
+            f"⏳ Job Started\nID: `{job_id}`\nTopic: _{topic}_"
         )
         
         script = generate_script(topic)
-        jobs[job_id] = {
-            "status": "complete",
-            "result": script
-        }
+        jobs[job_id] = {"status": "complete", "result": script}
         
         if callback_url:
             try:
                 async with httpx.AsyncClient() as client:
                     response = await client.post(
                         callback_url,
-                        json={
-                            "job_id": job_id,
-                            "status": "complete",
-                            "result": script
-                        },
+                        json={"job_id": job_id, "status": "complete", "result": script},
                         timeout=30.0
                     )
-                    
                     if response.status_code == 200:
                         await generator.send_telegram_notification(
-                            f"✅ *Job Completed*\n"
-                            f"ID: `{job_id}`\n"
-                            f"Callback successful"
-                        )
-                    else:
-                        error_msg = f"Callback failed with status {response.status_code}"
-                        logger.error(error_msg)
-                        await generator.send_telegram_notification(
-                            f"⚠️ *Callback Failed*\n"
-                            f"ID: `{job_id}`\n"
-                            f"Error: _{error_msg}_"
+                            f"✅ Job Complete\nID: `{job_id}`"
                         )
-                        
             except Exception as e:
-                error_msg = f"Callback error: {str(e)}"
-                logger.error(error_msg)
-                await generator.send_telegram_notification(
-                    f"⚠️ *Callback Error*\n"
-                    f"ID: `{job_id}`\n"
-                    f"Error: _{error_msg}_"
-                )
+                logger.error(f"Callback failed: {str(e)}")
         
     except Exception as e:
         error_msg = f"Job failed: {str(e)}"
-        logger.error(error_msg)
-        jobs[job_id] = {
-            "status": "failed",
-            "error": error_msg
-        }
-        
+        jobs[job_id] = {"status": "failed", "error": error_msg}
         await generator.send_telegram_notification(
-            f"❌ *Job Failed*\n"
-            f"ID: `{job_id}`\n"
-            f"Error: _{error_msg}_"
+            f"❌ Job Failed\nID: `{job_id}`\nError: _{error_msg}_"
         )
-        
-        if callback_url:
-            try:
-                await httpx.post(
-                    callback_url,
-                    json={
-                        "job_id": job_id,
-                        "status": "failed",
-                        "error": error_msg
-                    }
-                )
-            except Exception as e:
-                logger.error(f"Failed to send failure callback: {str(e)}")
 
 @app.post("/api/submit")
 async def submit_job(request: Request, background_tasks: BackgroundTasks):
-    """Endpoint to submit new job"""
+    """Submit new job"""
     try:
         data = await request.json()
         job_id = str(uuid.uuid4())
         
-        # Validate input
         if not data.get("topic"):
-            raise HTTPException(status_code=400, detail="Topic is required")
+            raise HTTPException(400, detail="Topic is required")
             
-        callback_url = data.get("callback_url")
-        
         jobs[job_id] = {
             "status": "processing",
-            "result": None,
-            "callback_url": callback_url
+            "callback_url": data.get("callback_url")
         }
         
         background_tasks.add_task(
             process_job,
             job_id,
             data["topic"],
-            callback_url
+            data.get("callback_url")
         )
         
-        return JSONResponse({
-            "job_id": job_id,
-            "status": "queued",
-            "message": "Job is being processed"
-        })
+        return {"job_id": job_id, "status": "queued"}
         
     except Exception as e:
         logger.error(f"Submission error: {str(e)}")
-        await generator.send_telegram_notification(
-            f"🚨 *Submission Error*\n"
-            f"Error: _{str(e)}_"
-        )
-        raise HTTPException(status_code=400, detail=str(e))
+        raise HTTPException(400, detail=str(e))
 
 @app.get("/api/status/{job_id}")
 async def get_status(job_id: str):
     """Check job status"""
     if job_id not in jobs:
-        raise HTTPException(status_code=404, detail="Job not found")
+        raise HTTPException(404, detail="Job not found")
     return jobs[job_id]
 
 @app.get("/health")
 async def health_check():
     """Health check endpoint"""
-    return {"status": "healthy", "model_loaded": generator.loaded}
+    return {
+        "status": "healthy",
+        "model_loaded": generator.loaded,
+        "dns_working": check_dns(),
+        "telegram_configured": bool(TELEGRAM_BOT_TOKEN and TELEGRAM_CHAT_ID)
+    }
 
 if __name__ == "__main__":
     uvicorn.run(