Update app.py

app.py

@@ -7,6 +7,7 @@ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
 import psycopg2
 from datetime import datetime
+from psycopg2 import pool
 
 app = Flask(__name__)
 CORS(app)
@@ -19,16 +20,33 @@ logger = logging.getLogger(__name__)
 MODEL_NAME = "tscholak/3vnuv1vf"  # نموذج متخصص لـ PostgreSQL
 SUPABASE_DB_URL = os.getenv('SUPABASE_DB_URL')
 
-# التهيئة المسبقة للنموذج
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
+# تهيئة connection pool
+connection_pool = None
+try:
+    connection_pool = psycopg2.pool.SimpleConnectionPool(
+        minconn=1,
+        maxconn=5,
+        dsn=SUPABASE_DB_URL
+    )
+    logger.info("تم إنشاء connection pool بنجاح")
+except Exception as e:
+    logger.error(f"خطأ في إنشاء connection pool: {str(e)}")
 
-# نقل النموذج إلى الجهاز المناسب
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model.to(device)
-model.eval()
+# تحميل النموذج مرة واحدة عند بدء التشغيل
+try:
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
+    
+    # استخدام GPU إذا كان متاحًا، وإلا استخدام CPU
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.to(device)
+    model.eval()
+    logger.info("تم تحميل النموذج بنجاح على الجهاز: %s", device)
+except Exception as e:
+    logger.error(f"خطأ في تحميل النموذج: {str(e)}")
+    raise
 
-# سكيما قاعدة البيانات
+# سكيما قاعدة البيانات (مختصرة لتحسين الأداء)
 DB_SCHEMA = """
 CREATE TABLE public.biodata (
   id bigint GENERATED ALWAYS AS IDENTITY NOT NULL,
@@ -110,30 +128,61 @@ CREATE TABLE public.user_place (
 );
 """.strip()
 
-def generate_sql(prompt: str, cam_mac: str) -> str:
+def clean_sql(sql: str, cam_mac: str) -> str:
     """
-    توليد استعلام SQL من النص باستخدام النموذج
+    تنظيف استعلام SQL وإضافة شروط الأمان
     """
-    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512).to(device)
+    # إزالة أي أوامر غير مسموح بها
+    forbidden_keywords = ['insert', 'update', 'delete', 'drop', 'alter', 'create', 'truncate']
+    for keyword in forbidden_keywords:
+        if keyword in sql.lower():
+            raise ValueError(f"استعلام غير مسموح به يحتوي على {keyword}")
     
-    with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_length=256,
-            num_beams=4,
-            early_stopping=True,
-            temperature=0.7
-        )
+    # التأكد من وجود شرط cam_mac
+    if 'where' in sql.lower():
+        sql = re.sub(r'where\s+', f"WHERE cam_mac = '{cam_mac}' AND ", sql, flags=re.IGNORECASE)
+    else:
+        if ';' in sql:
+            sql = sql.replace(';', f" WHERE cam_mac = '{cam_mac}';")
+        else:
+            sql += f" WHERE cam_mac = '{cam_mac}'"
     
-    sql = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return clean_sql(sql, cam_mac)
+    # التأكد من أن الاستعلام يبدأ بـ SELECT فقط
+    if not sql.strip().lower().startswith('select'):
+        raise ValueError("يسمح فقط باستعلامات SELECT")
+    
+    return sql
+
+def generate_sql(prompt: str, cam_mac: str) -> str:
+    """
+    توليد استعلام SQL من النص باستخدام النموذج
+    """
+    try:
+        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512).to(device)
+        
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_length=256,
+                num_beams=4,
+                early_stopping=True,
+                temperature=0.7
+            )
+        
+        sql = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return clean_sql(sql, cam_mac)
+    except Exception as e:
+        logger.error(f"خطأ في توليد SQL: {str(e)}")
+        return f"SELECT * FROM data WHERE cam_mac = '{cam_mac}' LIMIT 10;"
 
 @app.route('/api/query', methods=['POST'])
 def handle_query():
+    if not connection_pool:
+        return jsonify({"error": "لا يوجد اتصال بقاعدة البيانات"}), 500
+    
     try:
         data = request.get_json()
         
-        # التحقق من البيانات المدخلة
         if not all(k in data for k in ['text', 'cam_mac']):
             return jsonify({"error": "المعطيات ناقصة"}), 400
 
@@ -146,31 +195,45 @@ def handle_query():
 1. يجب تضمين شرط WHERE: cam_mac = '{data['cam_mac']}'
 2. مسموح فقط باستخدام SELECT
 3. الجداول المتاحة: profiles, data, place
-
-المثال:
-السؤال: "ما عدد زياراتي؟"
-SQL: SELECT COUNT(*) FROM data WHERE cam_mac = '{data['cam_mac']}';
 """
         # توليد الاستعلام
         sql = generate_sql(prompt, data['cam_mac'])
         
-        # التحقق من الأمان
-        if not is_safe_sql(sql):
-            sql = f"SELECT * FROM data WHERE cam_mac = '{data['cam_mac']}' LIMIT 10;"
-
-        # تنفيذ الاستعلام
-        with psycopg2.connect(SUPABASE_DB_URL) as conn:
+        # تنفيذ الاستعلام باستخدام connection pool
+        conn = None
+        try:
+            conn = connection_pool.getconn()
             with conn.cursor() as cursor:
                 cursor.execute(sql)
-                columns = [desc[0] for desc in cursor.description]
-                rows = cursor.fetchall()
+                
+                # إذا كان الاستعلام لا يعيد بيانات (مثل COUNT)
+                if cursor.description:
+                    columns = [desc[0] for desc in cursor.description]
+                    rows = cursor.fetchall()
+                    result = [dict(zip(columns, row)) for row in rows]
+                else:
+                    result = {"message": "تم تنفيذ الاستعلام بنجاح"}
                 
                 return jsonify({
-                    "data": [dict(zip(columns, row)) for row in rows],
+                    "data": result,
                     "sql": sql,
                     "timestamp": datetime.now().isoformat()
                 })
+        except Exception as e:
+            logger.error(f"خطأ في قاعدة البيانات: {str(e)}")
+            return jsonify({"error": "حدث خطأ في معالجة الاستعلام"}), 500
+        finally:
+            if conn:
+                connection_pool.putconn(conn)
                 
     except Exception as e:
-        logger.error(f"Error: {str(e)}")
-        return jsonify({"error": "حدث خطأ في المعالجة"}), 500
+        logger.error(f"خطأ عام: {str(e)}")
+        return jsonify({"error": "حدث خطأ في المعالجة"}), 500
+
+@app.route('/health', methods=['GET'])
+def health_check():
+    return jsonify({"status": "healthy", "timestamp": datetime.now().isoformat()})
+
+if __name__ == '__main__':
+    port = int(os.environ.get('PORT', 8080))
+    app.run(host='0.0.0.0', port=port)