Update main.py

main.py

@@ -1,4 +1,4 @@
-from fastapi import FastAPI, HTTPException, BackgroundTasks
+from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel, Field
 from sentence_transformers import SentenceTransformer
 import uvicorn
@@ -9,23 +9,19 @@ from contextlib import asynccontextmanager
 import httpx
 import os
 import sqlite3
-from datetime import datetime
 import json
-import threading
 
 # Globals
 model = None
 tokenizer = None
 model_id = 'Qwen/Qwen3-Embedding-0.6B'
 MAX_TOKENS = 32000
-DB_PATH = "/data/embeddings.db"  # هام: المسار ده في HuggingFace
-processing_lock = threading.Lock()
+DB_PATH = "/data/embeddings.db"
 is_processing = False
 
 
 def init_database():
     """Initialize the SQLite database"""
-    # Create /data directory if it doesn't exist
     os.makedirs(os.path.dirname(DB_PATH), exist_ok=True)
     
     conn = sqlite3.connect(DB_PATH)
@@ -45,7 +41,6 @@ def init_database():
         )
     ''')
     
-    # Create index for faster queries
     cursor.execute('''
         CREATE INDEX IF NOT EXISTS idx_status 
         ON embedding_requests(status)
@@ -53,7 +48,7 @@ def init_database():
     
     conn.commit()
     conn.close()
-    print("Database initialized successfully")
+    print("✅ Database initialized successfully")
 
 
 def save_request_to_db(text: str, request_id: str = None) -> int:
@@ -125,34 +120,44 @@ def update_embedding_in_db(row_id: int, embedding: List[float]):
     
     conn.commit()
     conn.close()
-    print(f"✅ Embedding updated for row ID: {row_id}")
+    print(f"✅ Embedding saved for ID: {row_id}")
 
 
-def mark_webhook_sent(row_id: int):
-    """Mark that webhook was sent successfully"""
+def get_request_data(row_id: int):
+    """Get full request data including embedding"""
     conn = sqlite3.connect(DB_PATH)
     cursor = conn.cursor()
     
     cursor.execute('''
-        UPDATE embedding_requests 
-        SET webhook_sent = 1
+        SELECT id, request_id, text, embedding
+        FROM embedding_requests 
         WHERE id = ?
     ''', (row_id,))
     
-    conn.commit()
+    result = cursor.fetchone()
     conn.close()
+    
+    return result
 
 
-def delete_from_db(row_id: int):
-    """Delete the request from database after webhook is sent"""
+def mark_webhook_sent_and_delete(row_id: int):
+    """Mark webhook as sent and delete from DB"""
     conn = sqlite3.connect(DB_PATH)
     cursor = conn.cursor()
     
+    # First mark as sent
+    cursor.execute('''
+        UPDATE embedding_requests 
+        SET webhook_sent = 1
+        WHERE id = ?
+    ''', (row_id,))
+    
+    # Then delete
     cursor.execute('DELETE FROM embedding_requests WHERE id = ?', (row_id,))
     
     conn.commit()
     conn.close()
-    print(f"🗑️ Request deleted from DB with ID: {row_id}")
+    print(f"🗑️ Request deleted from DB: {row_id}")
 
 
 def mark_request_failed(row_id: int, error_message: str):
@@ -177,42 +182,43 @@ async def lifespan(app: FastAPI):
     # Initialize database
     init_database()
     
-    # Load the model and tokenizer at startup
+    # Load the model
     global model, tokenizer
     print(f"Loading model: {model_id}...")
     model = SentenceTransformer(model_id)
     tokenizer = model.tokenizer
-    print("Model loaded successfully")
+    print("✅ Model loaded successfully")
     
-    # Start the background processor
+    # Start background processor
     asyncio.create_task(process_queue())
     
     yield
     
-    # Clean up
-    print("Cleaning up resources...")
+    # Cleanup
+    print("Cleaning up...")
     model = None
     tokenizer = None
 
 
 app = FastAPI(
-    title="Text Embedding API (Qwen/Qwen3-Embedding-0.6B)",
+    title="Text Embedding API with Queue",
     lifespan=lifespan
 )
 
 
 class TextRequest(BaseModel):
     text: str = Field(..., min_length=1, description="Text to embed")
-    request_id: str | None = Field(None, description="Optional unique identifier for the request")
+    request_id: str | None = Field(None, description="Optional request identifier")
 
 
 def chunk_and_embed(text: str) -> List[float]:
-    """Split text into chunks if too long, then pool embeddings"""
+    """Generate embedding with chunking if needed"""
     tokens = tokenizer.encode(text, add_special_tokens=False)
     
     if len(tokens) <= MAX_TOKENS:
         return model.encode(text, normalize_embeddings=True).tolist()
     
+    # Chunking
     chunks = []
     overlap = 50
     start = 0
@@ -232,74 +238,74 @@ def chunk_and_embed(text: str) -> List[float]:
     return final_embedding
 
 
-async def send_to_webhook(url: str, data: dict, db_row_id: int):
-    """Send data to webhook and delete from DB on success"""
+async def send_to_webhook(webhook_url: str, row_id: int, request_id: str, text: str, embedding: List[float]):
+    """Send complete data to webhook after embedding is ready"""
     try:
+        payload = {
+            "db_id": row_id,
+            "request_id": request_id,
+            "text": text,
+            "embedding": embedding,
+            "status": "completed"
+        }
+        
         async with httpx.AsyncClient(timeout=60.0) as client:
-            response = await client.post(url, json=data)
+            response = await client.post(webhook_url, json=payload)
             response.raise_for_status()
-            print(f"✅ Webhook sent successfully for ID: {db_row_id}")
+            print(f"✅ Webhook sent successfully for ID: {row_id}")
             
-            mark_webhook_sent(db_row_id)
-            delete_from_db(db_row_id)
+            # Delete from DB after successful webhook
+            mark_webhook_sent_and_delete(row_id)
             
     except Exception as e:
-        print(f"❌ Webhook error for ID {db_row_id}: {e}")
+        print(f"❌ Webhook error for ID {row_id}: {e}")
+        # Don't delete if webhook failed
 
 
 async def process_queue():
-    """Background task to process pending requests one by one"""
+    """Background processor - processes one request at a time"""
     global is_processing
     
     print("🚀 Queue processor started")
     
     while True:
         try:
-            # Check if there's a pending request
             pending = get_next_pending_request()
             
             if pending:
                 row_id, request_id, text = pending
-                
-                # Mark as processing
                 is_processing = True
                 update_request_processing(row_id)
                 
                 print(f"⚙️ Processing request ID: {row_id}")
                 
                 try:
-                    # Generate embedding (synchronous in async context)
+                    # Generate embedding
                     embedding = await asyncio.to_thread(chunk_and_embed, text)
                     
-                    # Save embedding to DB
+                    # Save to DB
                     update_embedding_in_db(row_id, embedding)
                     
-                    # Send to webhook if URL exists
+                    # Send to webhook with ALL data
                     webhook_url = os.environ.get("WEBHOOK_URL")
                     if webhook_url:
-                        payload = {
-                            "db_id": row_id,
-                            "text": text,
-                            "embedding": embedding,
-                            "request_id": request_id
-                        }
-                        await send_to_webhook(webhook_url, row_id, payload)
+                        await send_to_webhook(webhook_url, row_id, request_id, text, embedding)
                     else:
                         # No webhook, just delete
-                        delete_from_db(row_id)
+                        mark_webhook_sent_and_delete(row_id)
                     
                 except Exception as e:
-                    print(f"❌ Error processing request {row_id}: {e}")
+                    print(f"❌ Error processing {row_id}: {e}")
                     mark_request_failed(row_id, str(e))
                 
                 is_processing = False
                 
             else:
-                # No pending requests, wait a bit
+                # No pending requests
                 await asyncio.sleep(2)
                 
         except Exception as e:
-            print(f"❌ Queue processor error: {e}")
+            print(f"❌ Queue error: {e}")
             is_processing = False
             await asyncio.sleep(5)
 
@@ -308,24 +314,27 @@ async def process_queue():
 def home():
     return {
         "status": "online", 
-        "model": model_id, 
-        "endpoint": "/embed/text",
+        "model": model_id,
         "processing": is_processing
     }
 
 
 @app.post("/embed/text")
 async def embed_text(request: TextRequest):
-    """Just save the request to database, processing happens in background"""
+    """
+    Fast response - just queue the request
+    Processing happens in background
+    """
     try:
-        # Simply save to database
+        # Save to DB immediately
         db_row_id = save_request_to_db(request.text, request.request_id)
         
+        # Return immediately
         return {
             "success": True,
-            "message": "Request queued for processing",
+            "message": "Request queued successfully",
             "db_id": db_row_id,
-            "model": model_id
+            "status": "pending"
         }
         
     except Exception as e:
@@ -334,7 +343,7 @@ async def embed_text(request: TextRequest):
 
 @app.get("/status")
 def get_status():
-    """Get queue status"""
+    """Get queue statistics"""
     conn = sqlite3.connect(DB_PATH)
     cursor = conn.cursor()
     
@@ -350,24 +359,36 @@ def get_status():
     cursor.execute('SELECT COUNT(*) FROM embedding_requests WHERE status = "failed"')
     failed = cursor.fetchone()[0]
     
-    cursor.execute('SELECT COUNT(*) FROM embedding_requests')
-    total = cursor.fetchone()[0]
+    # Get next in queue
+    cursor.execute('''
+        SELECT id, created_at 
+        FROM embedding_requests 
+        WHERE status = "pending" 
+        ORDER BY id ASC 
+        LIMIT 1
+    ''')
+    next_request = cursor.fetchone()
     
     conn.close()
     
     return {
-        "total": total,
-        "pending": pending,
-        "processing": processing,
-        "completed": completed,
-        "failed": failed,
-        "is_processing": is_processing
+        "queue": {
+            "pending": pending,
+            "processing": processing,
+            "completed": completed,
+            "failed": failed
+        },
+        "is_processing": is_processing,
+        "next_request": {
+            "id": next_request[0] if next_request else None,
+            "created_at": next_request[1] if next_request else None
+        } if next_request else None
     }
 
 
 @app.get("/request/{db_id}")
-def get_request_status(db_id: int):
-    """Check status of a specific request"""
+def get_request_info(db_id: int):
+    """Check specific request status"""
     conn = sqlite3.connect(DB_PATH)
     cursor = conn.cursor()
     
@@ -381,7 +402,7 @@ def get_request_status(db_id: int):
     conn.close()
     
     if not result:
-        raise HTTPException(status_code=404, detail="Request not found")
+        raise HTTPException(status_code=404, detail="Request not found or already deleted")
     
     return {
         "db_id": result[0],