Update main.py

main.py

@@ -12,7 +12,9 @@ from google.genai import types
 
 app = FastAPI()
 
-# --- إعدادات المفاتيح ---
+# ============================================================
+# إعدادات المفاتيح
+# ============================================================
 TEXT_KEYS_RAW = os.environ.get("TEXT_API_KEYS", os.environ.get("GEMINI_API_KEY", ""))
 TEXT_API_KEYS = [k.strip() for k in TEXT_KEYS_RAW.split(",") if k.strip()]
 NVIDIA_API_KEY = os.environ.get("NVIDIA_API_KEY", "")
@@ -22,6 +24,9 @@ CURRENT_TEXT_KEY_INDEX = 0
 
 user_image_limits = defaultdict(lambda: {"hour": datetime.now().strftime("%Y-%m-%d %H"), "count": 0})
 
+# ============================================================
+# النماذج
+# ============================================================
 class FileData(BaseModel):
     mime_type: str
     data: str
@@ -32,15 +37,11 @@ class ChatRequest(BaseModel):
     message: str
     history: list
     files: list[FileData] = []
-    model: str = "flash"  # "flash" أو "pro"
+    model: str = "flash"
 
-def get_text_client():
-    global CURRENT_TEXT_KEY_INDEX
-    if not text_clients:
-        raise ValueError("لم يتم العثور على مفاتيح النصوص.")
-    return text_clients[CURRENT_TEXT_KEY_INDEX]
-
-# --- تعليمات النظام ---
+# ============================================================
+# تعليمات النظام
+# ============================================================
 SYSTEM_INSTRUCTION = """أنت Genisi، نموذج ذكاء اصطناعي متطور من مبادرة AnesNT.
 
 **معلومات عن AnesNT:**
@@ -54,221 +55,265 @@ SYSTEM_INSTRUCTION = """أنت Genisi، نموذج ذكاء اصطناعي مت
 3. استخدم البحث من Google عند الحاجة.
 4. استخدم الإيموجي بشكل مناسب."""
 
-@app.post("/chat")
-async def chat_endpoint(request: ChatRequest):
+# ============================================================
+# دوال مساعدة
+# ============================================================
+def get_text_client():
+    """الحصول على العميل الحالي مع تدوير المفاتيح عند الحاجة"""
     global CURRENT_TEXT_KEY_INDEX
-    msg_lower = request.message.strip().lower()
-    current_model = request.model
+    if not text_clients:
+        raise ValueError("لم يتم العثور على مفاتيح النصوص.")
+    return text_clients[CURRENT_TEXT_KEY_INDEX]
+
+def build_contents(history: list, message: str, files: list) -> tuple:
+    """بناء محتوى المحادثة للـ API"""
+    contents = []
+    has_images = False
     
-    # --- طلبات الصور (NVIDIA Flux) ---
-    image_triggers = ["ارسم", "صمم", "تخيل", "صورة ل", "draw", "generate", "imagine", "create", "عدل", "edit", "تصميم"]
-    is_image_request = any(msg_lower.startswith(trigger) for trigger in image_triggers)
+    # بناء التاريخ
+    for entry in history:
+        if entry.get('user') and str(entry['user']).strip():
+            contents.append(types.Content(role="user", parts=[types.Part.from_text(text=str(entry['user']).strip())]))
+        if entry.get('bot') and str(entry['bot']).strip():
+            bot_txt = entry['bot']
+            if '<div style="text-align:center;' in bot_txt:
+                bot_txt = "[صورة تم توليدها مسبقاً]"
+            contents.append(types.Content(role="model", parts=[types.Part.from_text(text=bot_txt.strip())]))
+    
+    # بناء أجزاء المستخدم
+    user_parts = []
+    if files:
+        for f in files:
+            try:
+                file_bytes = base64.b64decode(f.data)
+                user_parts.append(types.Part.from_bytes(data=file_bytes, mime_type=f.mime_type))
+                if f.mime_type.startswith('image/'):
+                    has_images = True
+            except Exception:
+                pass
+    
+    msg_text = message.strip()
+    if msg_text:
+        user_parts.append(types.Part.from_text(text=msg_text))
+    elif not msg_text and files:
+        user_parts.append(types.Part.from_text(text="يرجى تحليل المرفقات."))
+    else:
+        user_parts.append(types.Part.from_text(text="مرحبا"))
+    
+    contents.append(types.Content(role="user", parts=user_parts))
+    return contents, has_images
 
-    if is_image_request:
-        def generate_image_stream():
-            current_hour = datetime.now().strftime("%Y-%m-%d %H")
-            user_info = user_image_limits[request.user_id]
-            
-            if user_info.get("hour") != current_hour:
-                user_info["hour"] = current_hour
-                user_info["count"] = 0 
+def get_base_config(temperature: float, thinking_level: str, tools: list):
+    """الحصول على إعدادات النموذج الأساسية"""
+    return types.GenerateContentConfig(
+        temperature=temperature,
+        thinking_config=types.ThinkingConfig(thinking_level=thinking_level),
+        tools=tools,
+        system_instruction=[types.Part.from_text(text=SYSTEM_INSTRUCTION)]
+    )
 
-            if user_info["count"] >= 3:
-                yield "⚠️ **عذراً!** لقد استنفدت رصيدك الحالي لتوليد الصور (3 صور في الساعة). يرجى المحاولة لاحقاً! 🕒"
-                return
-            
-            if not NVIDIA_API_KEY:
-                yield "⚠️ **خطأ في السيرفر:** مفتاح `NVIDIA_API_KEY` غير موجود."
-                return
-            
-            try:
-                client = get_text_client()
-                translation_response = client.models.generate_content(
-                    model="gemma-4-31b-it", 
-                    contents=f"Translate this prompt to English for an image generator. Only return the English prompt: {request.message}"
-                )
-                final_prompt = translation_response.text.strip()
+# ============================================================
+# 🌟 دالة مستقلة لنموذج FLASH (سريع ومباشر)
+# ============================================================
+async def handle_flash_model(request: ChatRequest, client, contents: list, tools: list):
+    """نموذج Flash: استجابة سريعة ومباشرة"""
+    chosen_model = "gemma-4-31b-it"
+    config = get_base_config(temperature=0.7, thinking_level="MINIMAL", tools=tools)
+    
+    stream = client.models.generate_content_stream(
+        model=chosen_model,
+        contents=contents,
+        config=config
+    )
+    
+    def stream_generator():
+        try:
+            for chunk in stream:
+                if chunk.text:
+                    yield chunk.text
+        except Exception as e:
+            yield f"\n\n⚠️ **خطأ:** `{str(e)}`"
+    
+    return StreamingResponse(stream_generator(), media_type="text/plain")
 
-                invoke_url = "https://ai.api.nvidia.com/v1/genai/black-forest-labs/flux.2-klein-4b"
-                headers = {
-                    "Authorization": f"Bearer {NVIDIA_API_KEY}",
-                    "Accept": "application/json",
-                }
+# ============================================================
+# 🌟 دالة مستقلة لنموذج PRO (تفكير عميق + رد نهائي)
+# ============================================================
+async def handle_pro_model(request: ChatRequest, client, contents: list, tools: list):
+    """نموذج Pro: تفكير منفصل داخل حاوية، ثم رد نهائي خارجها"""
+    chosen_model = "gemma-4-31b-it"
+    
+    # ---------- المرحلة 1: توليد التفكير (بدون Streaming) ----------
+    thinking_config = get_base_config(temperature=1.0, thinking_level="HIGH", tools=tools)
+    
+    thinking_response = client.models.generate_content(
+        model=chosen_model,
+        contents=contents,
+        config=thinking_config
+    )
+    thinking_text = thinking_response.text.strip() if thinking_response.text else ""
+    
+    # ---------- المرحلة 2: توليد الرد النهائي (مع Streaming) ----------
+    # نضيف التفكير كجزء من السياق
+    final_prompt = f"""[لقد قمت بعملية تفكير عميق حول سؤال المستخدم، وهذا ما توصلت إليه:]
 
-                payload = {
-                    "prompt": final_prompt,
-                    "width": 1024,
-                    "height": 1024,
-                    "seed": 0,
-                    "steps": 4
-                }
+{thinking_text}
 
-                if request.files:
-                    img_file = next((f for f in request.files if f.mime_type.startswith("image/")), None)
-                    if img_file:
-                        payload["image"] = [f"data:{img_file.mime_type};base64,{img_file.data}"]
+[الآن، بناءً على هذا التفكير، قدم رداً نهائياً مباشراً ومفيداً للمستخدم. لا تذكر أنك فكرت، فقط قدم الرد النهائي.]"""
+    
+    final_contents = contents.copy()
+    final_contents.append(types.Content(role="user", parts=[types.Part.from_text(text=final_prompt)]))
+    
+    final_config = get_base_config(temperature=0.7, thinking_level="MINIMAL", tools=tools)
+    
+    stream = client.models.generate_content_stream(
+        model=chosen_model,
+        contents=final_contents,
+        config=final_config
+    )
+    
+    def stream_generator():
+        # نرسل التفكير + علامة الفصل أولاً
+        yield thinking_text + " instant"
+        
+        # ثم نرسل الرد النهائي
+        try:
+            for chunk in stream:
+                if chunk.text:
+                    yield chunk.text
+        except Exception as e:
+            yield f"\n\n⚠️ **خطأ:** `{str(e)}`"
+    
+    return StreamingResponse(stream_generator(), media_type="text/plain")
 
-                response = requests.post(invoke_url, headers=headers, json=payload, timeout=60)
-                
-                if not response.ok:
-                    yield f"⚠️ **خطأ من سيرفر NVIDIA:**\n`{response.text}`"
-                    return
+# ============================================================
+# 🌟 دالة مستقلة لتوليد الصور
+# ============================================================
+async def handle_image_request(request: ChatRequest, client):
+    """معالجة طلبات توليد الصور باستخدام NVIDIA Flux"""
+    def generate_image_stream():
+        current_hour = datetime.now().strftime("%Y-%m-%d %H")
+        user_info = user_image_limits[request.user_id]
+        
+        if user_info.get("hour") != current_hour:
+            user_info["hour"] = current_hour
+            user_info["count"] = 0 
 
-                response_body = response.json()
-                
-                def extract_img(obj):
-                    if isinstance(obj, dict):
-                        for k, v in obj.items():
-                            if k in ['b64_json', 'base64', 'image'] and isinstance(v, str) and len(v) > 100:
-                                return v.split(",", 1)[-1] if v.startswith("data:") else v
-                            elif k == 'url' and isinstance(v, str) and v.startswith('http'):
-                                return base64.b64encode(requests.get(v).content).decode('utf-8')
-                            else:
-                                res = extract_img(v)
-                                if res: return res
-                    elif isinstance(obj, list):
-                        for item in obj:
-                            res = extract_img(item)
+        if user_info["count"] >= 3:
+            yield "⚠️ **عذراً!** لقد استنفدت رصيدك الحالي لتوليد الصور (3 صور في الساعة). يرجى المحاولة لاحقاً! 🕒"
+            return
+        
+        if not NVIDIA_API_KEY:
+            yield "⚠️ **خطأ في السيرفر:** مفتاح `NVIDIA_API_KEY` غير موجود."
+            return
+        
+        try:
+            # ترجمة الطلب للإنجليزية
+            translation_response = client.models.generate_content(
+                model="gemma-4-31b-it", 
+                contents=f"Translate this prompt to English for an image generator. Only return the English prompt: {request.message}"
+            )
+            final_prompt = translation_response.text.strip()
+
+            invoke_url = "https://ai.api.nvidia.com/v1/genai/black-forest-labs/flux.2-klein-4b"
+            headers = {
+                "Authorization": f"Bearer {NVIDIA_API_KEY}",
+                "Accept": "application/json",
+            }
+
+            payload = {
+                "prompt": final_prompt,
+                "width": 1024,
+                "height": 1024,
+                "seed": 0,
+                "steps": 4
+            }
+
+            if request.files:
+                img_file = next((f for f in request.files if f.mime_type.startswith("image/")), None)
+                if img_file:
+                    payload["image"] = [f"data:{img_file.mime_type};base64,{img_file.data}"]
+
+            response = requests.post(invoke_url, headers=headers, json=payload, timeout=60)
+            
+            if not response.ok:
+                yield f"⚠️ **خطأ من سيرفر NVIDIA:**\n`{response.text}`"
+                return
+
+            response_body = response.json()
+            
+            def extract_img(obj):
+                if isinstance(obj, dict):
+                    for k, v in obj.items():
+                        if k in ['b64_json', 'base64', 'image'] and isinstance(v, str) and len(v) > 100:
+                            return v.split(",", 1)[-1] if v.startswith("data:") else v
+                        elif k == 'url' and isinstance(v, str) and v.startswith('http'):
+                            return base64.b64encode(requests.get(v).content).decode('utf-8')
+                        else:
+                            res = extract_img(v)
                             if res: return res
-                    return None
-                    
-                b64_image = extract_img(response_body)
+                elif isinstance(obj, list):
+                    for item in obj:
+                        res = extract_img(item)
+                        if res: return res
+                return None
                 
-                if not b64_image:
-                     yield f"⚠️ **رد غير متوقع من الخادم (لم يتم العثور على الصورة):**"
-                     return
+            b64_image = extract_img(response_body)
+            
+            if not b64_image:
+                 yield f"⚠️ **رد غير متوقع من الخادم (لم يتم العثور على الصورة):**"
+                 return
 
-                user_info["count"] += 1
-                rem = 3 - user_info["count"]
-                
-                html_response = f'''
+            user_info["count"] += 1
+            rem = 3 - user_info["count"]
+            
+            html_response = f'''
 <div style="text-align:center; margin: 15px 0;">
     <img src="data:image/jpeg;base64,{b64_image}" style="width:100%; max-width:400px; border-radius:18px; box-shadow:0 8px 25px rgba(0,0,0,0.15);" />
     <br/>
     <a href="data:image/jpeg;base64,{b64_image}" download="Genisi_Flux_Art.jpg" style="display:inline-block; margin-top:12px; padding:10px 20px; background:linear-gradient(135deg, #4f8ef7, #7c5cf7); color:#fff; border-radius:25px; text-decoration:none; font-weight:600; font-family:'Cairo', sans-serif;">⬇️ تحميل الصورة</a>
     <p style="font-size:0.85rem; color:#9aa3be; margin-top:8px;">✅ تم التصميم بنجاح (المتبقي لك هذه الساعة: {rem}/3 صور)</p>
 </div>'''
-                yield html_response
-            except Exception as e:
-                yield f"⚠️ **خطأ أثناء توليد الصورة:**\n`{str(e)}`"
-        
-        return StreamingResponse(generate_image_stream(), media_type="text/plain")
+            yield html_response
+        except Exception as e:
+            yield f"⚠️ **خطأ أثناء توليد الصورة:**\n`{str(e)}`"
+    
+    return StreamingResponse(generate_image_stream(), media_type="text/plain")
+
+# ============================================================
+# نقطة النهاية الرئيسية
+# ============================================================
+@app.post("/chat")
+async def chat_endpoint(request: ChatRequest):
+    global CURRENT_TEXT_KEY_INDEX
+    
+    msg_lower = request.message.strip().lower()
+    current_model = request.model
+    
+    # التحقق من طلبات الصور
+    image_triggers = ["ارسم", "صمم", "تخيل", "صورة ل", "draw", "generate", "imagine", "create", "عدل", "edit", "تصميم"]
+    is_image_request = any(msg_lower.startswith(trigger) for trigger in image_triggers)
 
-    # --- توليد النصوص ---
+    # محاولة تنفيذ الطلب مع تدوير المفاتيح
     attempts = 0
     while attempts < len(text_clients):
         try:
             client = get_text_client()
-            contents = []
-            
-            # بناء التاريخ
-            for entry in request.history:
-                if entry.get('user') and str(entry['user']).strip():
-                    contents.append(types.Content(role="user", parts=[types.Part.from_text(text=str(entry['user']).strip())]))
-                if entry.get('bot') and str(entry['bot']).strip():
-                    bot_txt = entry['bot']
-                    if '<div style="text-align:center;' in bot_txt:
-                        bot_txt = "[صورة تم توليدها مسبقاً]"
-                    contents.append(types.Content(role="model", parts=[types.Part.from_text(text=bot_txt.strip())]))
-            
-            # أجزاء المستخدم
-            user_parts = []
-            has_images = False
             
-            if request.files:
-                for f in request.files:
-                    try:
-                        file_bytes = base64.b64decode(f.data)
-                        user_parts.append(types.Part.from_bytes(data=file_bytes, mime_type=f.mime_type))
-                        if f.mime_type.startswith('image/'):
-                            has_images = True
-                    except Exception:
-                        pass
-            
-            msg_text = request.message.strip()
-            if msg_text:
-                user_parts.append(types.Part.from_text(text=msg_text))
-            elif not msg_text and request.files:
-                user_parts.append(types.Part.from_text(text="يرجى تحليل المرفقات."))
-            else:
-                user_parts.append(types.Part.from_text(text="مرحبا"))
-
-            contents.append(types.Content(role="user", parts=user_parts))
+            # --- طلب صورة ---
+            if is_image_request:
+                return await handle_image_request(request, client)
             
+            # --- طلب نصي ---
+            contents, _ = build_contents(request.history, request.message, request.files)
             tools = [types.Tool(googleSearch=types.GoogleSearch())]
-            chosen_model = "gemma-4-31b-it"
-
-            # --- 🌟 منطق Pro: تفكير منفصل ثم رد نهائي مع علامة " instant" ---
-            if current_model == "pro":
-                # المرحلة 1: توليد التفكير (بدون streaming)
-                thinking_config = types.GenerateContentConfig(
-                    temperature=1.0,
-                    thinking_config=types.ThinkingConfig(thinking_level="HIGH"),
-                    tools=tools,
-                    system_instruction=[types.Part.from_text(text=SYSTEM_INSTRUCTION)]
-                )
-                thinking_response = client.models.generate_content(
-                    model=chosen_model,
-                    contents=contents,
-                    config=thinking_config
-                )
-                thinking_text = thinking_response.text.strip()
-                
-                # المرحلة 2: توليد الرد النهائي (مع streaming)
-                # نضيف التفكير كجزء من السياق عشان النموذج يعرف إنو خلاص فكر
-                final_prompt = f"[لقد فكرت مسبقاً بما يلي]\n{thinking_text}\n\n[الآن قدم ردك النهائي المباشر بناءً على هذا التفكير]"
-                # نضيفه كجزء من محتوى المستخدم الأخير (أو نضيف Content جديد)
-                final_contents = contents.copy()
-                final_contents.append(types.Content(role="user", parts=[types.Part.from_text(text=final_prompt)]))
-                
-                final_config = types.GenerateContentConfig(
-                    temperature=0.7,
-                    thinking_config=types.ThinkingConfig(thinking_level="MINIMAL"),  # بدون تفكير إضافي
-                    tools=tools,
-                    system_instruction=[types.Part.from_text(text=SYSTEM_INSTRUCTION)]
-                )
-                
-                stream = client.models.generate_content_stream(
-                    model=chosen_model,
-                    contents=final_contents,
-                    config=final_config
-                )
-                
-                def stream_generator():
-                    # أولاً نرسل التفكير + علامة الفصل
-                    yield thinking_text + " instant"
-                    # ثم نرسل الرد النهائي
-                    try:
-                        for chunk in stream:
-                            if chunk.text:
-                                yield chunk.text
-                    except Exception as e:
-                        yield f"\n\n⚠️ **خطأ:** `{str(e)}`"
-                
-                return StreamingResponse(stream_generator(), media_type="text/plain")
             
-            # --- 🌟 منطق Flash: عادي مباشر ---
+            # 🌟 توجيه إلى الدالة المناسبة حسب النموذج
+            if current_model == "pro":
+                return await handle_pro_model(request, client, contents, tools)
             else:
-                config = types.GenerateContentConfig(
-                    temperature=0.7,
-                    thinking_config=types.ThinkingConfig(thinking_level="MINIMAL"),
-                    tools=tools,
-                    system_instruction=[types.Part.from_text(text=SYSTEM_INSTRUCTION)]
-                )
-                stream = client.models.generate_content_stream(
-                    model=chosen_model,
-                    contents=contents,
-                    config=config
-                )
-                def stream_generator():
-                    try:
-                        for chunk in stream:
-                            if chunk.text:
-                                yield chunk.text
-                    except Exception as e:
-                        yield f"\n\n⚠️ **خطأ:** `{str(e)}`"
-                return StreamingResponse(stream_generator(), media_type="text/plain")
-
+                return await handle_flash_model(request, client, contents, tools)
+                
         except Exception as e:
             error_msg = str(e).lower()
             if "429" in error_msg or "quota" in error_msg:
@@ -277,8 +322,15 @@ async def chat_endpoint(request: ChatRequest):
             else:
                 def err_gen(): yield f"⚠️ **خطأ في النموذج:** {str(e)}"
                 return StreamingResponse(err_gen(), media_type="text/plain")
-
+    
     def limit_gen(): yield "⚠️ تم الوصول للحد الأقصى لجميع المفاتيح."
     return StreamingResponse(limit_gen(), media_type="text/plain")
 
-app.mount("/", StaticFiles(directory=".", html=True), name="static")
+# ============================================================
+# تقديم الملفات الثابتة
+# ============================================================
+app.mount("/", StaticFiles(directory=".", html=True), name="static")
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)