Update app.py

app.py

@@ -3,21 +3,18 @@ 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
 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"
 HF_TOKEN = os.getenv("HF_TOKEN", "")
-TELEGRAM_BOT_TOKEN = os.getenv("TELEGRAM_BOT_TOKEN")
-TELEGRAM_CHAT_ID = os.getenv("TELEGRAM_CHAT_ID")
+TELEGRAM_BOT_TOKEN = os.getenv("TELEGRAM_BOT_TOKEN", "")
+TELEGRAM_CHAT_ID = os.getenv("TELEGRAM_CHAT_ID", "")
 MAX_TOKENS = 400
 DEVICE = "cpu"
 PORT = int(os.getenv("PORT", 7860))
@@ -32,10 +29,6 @@ logger = logging.getLogger(__name__)
 # Job storage
 jobs: Dict[str, dict] = {}
 
-class JobRequest(BaseModel):
-    topic: str
-    callback_url: Optional[str] = None
-
 class ScriptGenerator:
     def __init__(self):
         self.tokenizer = None
@@ -58,18 +51,14 @@ class ScriptGenerator:
             ).to(DEVICE)
             self.loaded = True
             logger.info("Model loaded successfully")
-            self.send_telegram_notification("🤖 Model loaded and ready!")
+            self._send_telegram_sync("🤖 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}")
+            logger.error(f"Model loading failed: {str(e)}")
             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 with retry logic"""
+    def _send_telegram_sync(self, message: str):
+        """Synchronous Telegram notification with retry"""
         if not TELEGRAM_BOT_TOKEN or not TELEGRAM_CHAT_ID:
-            logger.warning("Telegram not configured")
             return
             
         try:
@@ -79,20 +68,11 @@ class ScriptGenerator:
                 "text": message,
                 "parse_mode": "Markdown"
             }
-            
-            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")
-                
+            response = httpx.post(url, json=payload, timeout=10.0)
+            response.raise_for_status()
         except Exception as e:
-            logger.error(f"Telegram notification failed: {str(e)}")
-            raise
+            logger.warning(f"Telegram notification failed: {str(e)}")
 
-# Lifespan handler
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     generator = ScriptGenerator()
@@ -102,14 +82,6 @@ async def lifespan(app: FastAPI):
 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:
@@ -145,69 +117,72 @@ def generate_script(topic: str) -> str:
 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\nID: `{job_id}`\nTopic: _{topic}_"
-        )
+        generator._send_telegram_sync(f"⏳ Starting job: {job_id}\nTopic: {topic[:50]}...")
         
         script = generate_script(topic)
-        jobs[job_id] = {"status": "complete", "result": script}
+        jobs[job_id] = {
+            "status": "complete",
+            "result": script
+        }
+        
+        generator._send_telegram_sync(f"✅ Completed job: {job_id}")
         
         if callback_url:
             try:
-                async with httpx.AsyncClient() as client:
-                    response = await client.post(
+                async with httpx.AsyncClient(timeout=30.0) as client:
+                    await client.post(
                         callback_url,
-                        json={"job_id": job_id, "status": "complete", "result": script},
-                        timeout=30.0
+                        json={
+                            "job_id": job_id,
+                            "status": "complete",
+                            "result": script
+                        }
                     )
-                    if response.status_code == 200:
-                        await generator.send_telegram_notification(
-                            f"✅ Job Complete\nID: `{job_id}`"
-                        )
             except Exception as e:
                 logger.error(f"Callback failed: {str(e)}")
+                generator._send_telegram_sync(f"⚠️ Callback failed for {job_id}")
         
     except Exception as e:
         error_msg = f"Job failed: {str(e)}"
-        jobs[job_id] = {"status": "failed", "error": error_msg}
-        await generator.send_telegram_notification(
-            f"❌ Job Failed\nID: `{job_id}`\nError: _{error_msg}_"
-        )
+        jobs[job_id] = {
+            "status": "failed",
+            "error": error_msg
+        }
+        generator._send_telegram_sync(f"❌ Failed job: {job_id}\nError: {error_msg[:100]}")
 
 @app.post("/api/submit")
 async def submit_job(request: Request, background_tasks: BackgroundTasks):
-    """Submit new job"""
+    """Endpoint to submit new job"""
     try:
         data = await request.json()
         job_id = str(uuid.uuid4())
         
         if not data.get("topic"):
-            raise HTTPException(400, detail="Topic is required")
+            raise HTTPException(status_code=400, detail="Topic is required")
             
+        callback_url = data.get("callback_url")
+        
         jobs[job_id] = {
             "status": "processing",
-            "callback_url": data.get("callback_url")
+            "result": None,
+            "callback_url": callback_url
         }
         
         background_tasks.add_task(
             process_job,
             job_id,
             data["topic"],
-            data.get("callback_url")
+            callback_url
         )
         
-        return {"job_id": job_id, "status": "queued"}
+        return JSONResponse({
+            "job_id": job_id,
+            "status": "queued"
+        })
         
     except Exception as e:
         logger.error(f"Submission error: {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(404, detail="Job not found")
-    return jobs[job_id]
+        raise HTTPException(status_code=400, detail=str(e))
 
 @app.get("/health")
 async def health_check():
@@ -215,7 +190,6 @@ async def health_check():
     return {
         "status": "healthy",
         "model_loaded": generator.loaded,
-        "dns_working": check_dns(),
         "telegram_configured": bool(TELEGRAM_BOT_TOKEN and TELEGRAM_CHAT_ID)
     }
 
@@ -224,6 +198,5 @@ if __name__ == "__main__":
         app,
         host="0.0.0.0",
         port=PORT,
-        log_level="info",
-        timeout_keep_alive=300
+        log_level="info"
     )