Update main.py

main.py

@@ -3,69 +3,216 @@ from pydantic import BaseModel, Field
 from sentence_transformers import SentenceTransformer
 import uvicorn
 import asyncio
-from concurrent.futures import ThreadPoolExecutor
 from typing import List
 import numpy as np
 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'
-executor = ThreadPoolExecutor(max_workers=4)
 MAX_TOKENS = 32000
+DB_PATH = "/data/embeddings.db"  # هام: المسار ده في HuggingFace
+processing_lock = threading.Lock()
+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)
+    cursor = conn.cursor()
+    
+    cursor.execute('''
+        CREATE TABLE IF NOT EXISTS embedding_requests (
+            id INTEGER PRIMARY KEY AUTOINCREMENT,
+            request_id TEXT,
+            text TEXT NOT NULL,
+            embedding TEXT,
+            status TEXT DEFAULT 'pending',
+            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+            processed_at TIMESTAMP,
+            webhook_sent BOOLEAN DEFAULT 0,
+            error_message TEXT
+        )
+    ''')
+    
+    # Create index for faster queries
+    cursor.execute('''
+        CREATE INDEX IF NOT EXISTS idx_status 
+        ON embedding_requests(status)
+    ''')
+    
+    conn.commit()
+    conn.close()
+    print("Database initialized successfully")
+
+
+def save_request_to_db(text: str, request_id: str = None) -> int:
+    """Save the incoming request to database"""
+    conn = sqlite3.connect(DB_PATH)
+    cursor = conn.cursor()
+    
+    cursor.execute('''
+        INSERT INTO embedding_requests (request_id, text, status)
+        VALUES (?, ?, 'pending')
+    ''', (request_id, text))
+    
+    row_id = cursor.lastrowid
+    conn.commit()
+    conn.close()
+    
+    print(f"✅ Request saved to DB with ID: {row_id}")
+    return row_id
+
+
+def get_next_pending_request():
+    """Get the next pending request from database"""
+    conn = sqlite3.connect(DB_PATH)
+    cursor = conn.cursor()
+    
+    cursor.execute('''
+        SELECT id, request_id, text 
+        FROM embedding_requests 
+        WHERE status = 'pending'
+        ORDER BY id ASC
+        LIMIT 1
+    ''')
+    
+    result = cursor.fetchone()
+    conn.close()
+    
+    return result
+
+
+def update_request_processing(row_id: int):
+    """Mark request as processing"""
+    conn = sqlite3.connect(DB_PATH)
+    cursor = conn.cursor()
+    
+    cursor.execute('''
+        UPDATE embedding_requests 
+        SET status = 'processing'
+        WHERE id = ?
+    ''', (row_id,))
+    
+    conn.commit()
+    conn.close()
+
+
+def update_embedding_in_db(row_id: int, embedding: List[float]):
+    """Update the embedding in database"""
+    conn = sqlite3.connect(DB_PATH)
+    cursor = conn.cursor()
+    
+    embedding_json = json.dumps(embedding)
+    
+    cursor.execute('''
+        UPDATE embedding_requests 
+        SET embedding = ?, 
+            status = 'completed',
+            processed_at = CURRENT_TIMESTAMP
+        WHERE id = ?
+    ''', (embedding_json, row_id))
+    
+    conn.commit()
+    conn.close()
+    print(f"✅ Embedding updated for row ID: {row_id}")
+
+
+def mark_webhook_sent(row_id: int):
+    """Mark that webhook was sent successfully"""
+    conn = sqlite3.connect(DB_PATH)
+    cursor = conn.cursor()
+    
+    cursor.execute('''
+        UPDATE embedding_requests 
+        SET webhook_sent = 1
+        WHERE id = ?
+    ''', (row_id,))
+    
+    conn.commit()
+    conn.close()
+
+
+def delete_from_db(row_id: int):
+    """Delete the request from database after webhook is sent"""
+    conn = sqlite3.connect(DB_PATH)
+    cursor = conn.cursor()
+    
+    cursor.execute('DELETE FROM embedding_requests WHERE id = ?', (row_id,))
+    
+    conn.commit()
+    conn.close()
+    print(f"🗑️ Request deleted from DB with ID: {row_id}")
+
+
+def mark_request_failed(row_id: int, error_message: str):
+    """Mark request as failed"""
+    conn = sqlite3.connect(DB_PATH)
+    cursor = conn.cursor()
+    
+    cursor.execute('''
+        UPDATE embedding_requests 
+        SET status = 'failed',
+            error_message = ?,
+            processed_at = CURRENT_TIMESTAMP
+        WHERE id = ?
+    ''', (error_message, row_id))
+    
+    conn.commit()
+    conn.close()
+
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
+    # Initialize database
+    init_database()
+    
     # Load the model and tokenizer at startup
     global model, tokenizer
     print(f"Loading model: {model_id}...")
     model = SentenceTransformer(model_id)
     tokenizer = model.tokenizer
     print("Model loaded successfully")
+    
+    # Start the background processor
+    asyncio.create_task(process_queue())
+    
     yield
-    # (Optional) Clean up resources at shutdown
+    
+    # Clean up
     print("Cleaning up resources...")
     model = None
     tokenizer = None
 
+
 app = FastAPI(
     title="Text Embedding API (Qwen/Qwen3-Embedding-0.6B)",
     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")
 
 
-
-
-async def send_to_webhook(url: str, data: dict):
-    """Sends data to a webhook URL asynchronously."""
-    try:
-        async with httpx.AsyncClient() as client:
-            response = await client.post(url, json=data)
-            response.raise_for_status()  # Raise an exception for bad status codes (4xx or 5xx)
-            print(f"Successfully sent data to webhook: {url}")
-    except httpx.RequestError as e:
-        print(f"Error sending data to webhook {url}: {e}")
-
-@app.get("/")
-def home():
-    return {"status": "online", "model": model_id, "endpoint": "/embed/text"}
-
 def chunk_and_embed(text: str) -> List[float]:
     """Split text into chunks if too long, then pool embeddings"""
     tokens = tokenizer.encode(text, add_special_tokens=False)
     
-    # If text is short, embed directly
     if len(tokens) <= MAX_TOKENS:
         return model.encode(text, normalize_embeddings=True).tolist()
     
-    # Split into chunks
     chunks = []
     overlap = 50
     start = 0
@@ -79,39 +226,173 @@ def chunk_and_embed(text: str) -> List[float]:
             break
         start = end - overlap
     
-    # Embed all chunks
     chunk_embeddings = [model.encode(chunk, normalize_embeddings=True) for chunk in chunks]
-    
-    # Pool embeddings (mean)
     final_embedding = np.mean(chunk_embeddings, axis=0).tolist()
     
     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"""
+    try:
+        async with httpx.AsyncClient(timeout=60.0) as client:
+            response = await client.post(url, json=data)
+            response.raise_for_status()
+            print(f"✅ Webhook sent successfully for ID: {db_row_id}")
+            
+            mark_webhook_sent(db_row_id)
+            delete_from_db(db_row_id)
+            
+    except Exception as e:
+        print(f"❌ Webhook error for ID {db_row_id}: {e}")
+
+
+async def process_queue():
+    """Background task to process pending requests one by one"""
+    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)
+                    embedding = await asyncio.to_thread(chunk_and_embed, text)
+                    
+                    # Save embedding to DB
+                    update_embedding_in_db(row_id, embedding)
+                    
+                    # Send to webhook if URL exists
+                    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)
+                    else:
+                        # No webhook, just delete
+                        delete_from_db(row_id)
+                    
+                except Exception as e:
+                    print(f"❌ Error processing request {row_id}: {e}")
+                    mark_request_failed(row_id, str(e))
+                
+                is_processing = False
+                
+            else:
+                # No pending requests, wait a bit
+                await asyncio.sleep(2)
+                
+        except Exception as e:
+            print(f"❌ Queue processor error: {e}")
+            is_processing = False
+            await asyncio.sleep(5)
+
+
+@app.get("/")
+def home():
+    return {
+        "status": "online", 
+        "model": model_id, 
+        "endpoint": "/embed/text",
+        "processing": is_processing
+    }
+
+
 @app.post("/embed/text")
-async def embed_text(request: TextRequest, background_tasks: BackgroundTasks):
+async def embed_text(request: TextRequest):
+    """Just save the request to database, processing happens in background"""
     try:
-        loop = asyncio.get_event_loop()
-        embedding = await loop.run_in_executor(
-            executor,
-            lambda: chunk_and_embed(request.text)
-        )
+        # Simply save to database
+        db_row_id = save_request_to_db(request.text, request.request_id)
         
-        # Check for webhook URL and add the background task
-        webhook_url = os.environ.get("WEBHOOK_URL")
-        if webhook_url:
-            payload = {
-                "text": request.text,
-                "embedding": embedding,
-                "request_id": request.request_id 
-            }
-            background_tasks.add_task(send_to_webhook, webhook_url, payload)
-
         return {
             "success": True,
-            "model": model_id,
+            "message": "Request queued for processing",
+            "db_id": db_row_id,
+            "model": model_id
         }
+        
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
+
+@app.get("/status")
+def get_status():
+    """Get queue status"""
+    conn = sqlite3.connect(DB_PATH)
+    cursor = conn.cursor()
+    
+    cursor.execute('SELECT COUNT(*) FROM embedding_requests WHERE status = "pending"')
+    pending = cursor.fetchone()[0]
+    
+    cursor.execute('SELECT COUNT(*) FROM embedding_requests WHERE status = "processing"')
+    processing = cursor.fetchone()[0]
+    
+    cursor.execute('SELECT COUNT(*) FROM embedding_requests WHERE status = "completed"')
+    completed = cursor.fetchone()[0]
+    
+    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]
+    
+    conn.close()
+    
+    return {
+        "total": total,
+        "pending": pending,
+        "processing": processing,
+        "completed": completed,
+        "failed": failed,
+        "is_processing": is_processing
+    }
+
+
+@app.get("/request/{db_id}")
+def get_request_status(db_id: int):
+    """Check status of a specific request"""
+    conn = sqlite3.connect(DB_PATH)
+    cursor = conn.cursor()
+    
+    cursor.execute('''
+        SELECT id, request_id, status, created_at, processed_at, webhook_sent, error_message
+        FROM embedding_requests 
+        WHERE id = ?
+    ''', (db_id,))
+    
+    result = cursor.fetchone()
+    conn.close()
+    
+    if not result:
+        raise HTTPException(status_code=404, detail="Request not found")
+    
+    return {
+        "db_id": result[0],
+        "request_id": result[1],
+        "status": result[2],
+        "created_at": result[3],
+        "processed_at": result[4],
+        "webhook_sent": bool(result[5]),
+        "error_message": result[6]
+    }
+
+
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)