Update app.py

app.py

@@ -19,13 +19,16 @@ HF_TOKEN = os.environ.get("HF_TOKEN")
 TG_BOT_TOKEN = os.environ.get("TELEGRAM_BOT_TOKEN") 
 TG_CHAT_ID = os.environ.get("TELEGRAM_CHAT_ID")     
 
-# MODELS
+# MODELS (Router friendly models)
 TEXT_MODEL = "Qwen/Qwen2.5-Coder-32B-Instruct"
-VISION_MODEL = "Qwen/Qwen2-VL-7B-Instruct" # Sabse fast vision model
+VISION_MODEL = "meta-llama/Llama-3.2-11B-Vision-Instruct"
+
+# GLOBAL ROUTER URL (Jo HF ne manga hai)
+ROUTER_URL = "https://router.huggingface.co/v1/chat/completions"
 
 @app.route('/', methods=['GET'])
 def home():
-    return jsonify({"status": "TS AI Brain (Final Debug Mode) is Active! 🔥"})
+    return jsonify({"status": "TS AI Brain (Official Router Mode) is Active! 🔥"})
 
 @app.route('/api/chat', methods=['POST'])
 def chat():
@@ -34,10 +37,10 @@ def chat():
     image_url = None
 
     if is_form_data or request.files:
-        user_message = request.form.get("message", "Is photo ko read karo.")
+        user_message = request.form.get("message", "Explain this image.")
         image_file = request.files.get("image")
         if image_file:
-            # 1. TELEGRAM UPLOAD
+            # 1. TELEGRAM UPLOAD (STORAGE)
             tg_url = f"https://api.telegram.org/bot{TG_BOT_TOKEN}/sendPhoto"
             files = {'photo': (image_file.filename, image_file.read(), image_file.mimetype)}
             data = {'chat_id': TG_CHAT_ID}
@@ -49,6 +52,7 @@ def chat():
                     image_url = f"https://api.telegram.org/file/bot{TG_BOT_TOKEN}/{file_info['result']['file_path']}"
             except Exception as e:
                 return jsonify({"error": "Telegram upload fail", "details": str(e)}), 500
+
     else:
         data = request.json
         user_message = data.get("message", "")
@@ -58,46 +62,50 @@ def chat():
     doc = doc_ref.get()
     history = doc.to_dict().get('messages', []) if doc.exists else [{"role": "system", "content": "Tum TS Boss ke AI ho. Hinglish mein baat karo."}]
 
-    headers = {"Authorization": f"Bearer {HF_TOKEN}", "Content-Type": "application/json"}
+    headers = {
+        "Authorization": f"Bearer {HF_TOKEN}",
+        "Content-Type": "application/json",
+        "x-wait-for-model": "true" # Ye HF ko bolta hai ki model load hone tak ruko!
+    }
 
     try:
         if image_url:
-            # 2. VISION API CALL (Direct Model Path)
+            # --- VISION CALL VIA ROUTER ---
             payload = {
                 "model": VISION_MODEL,
-                "messages": [{"role": "user", "content": [{"type": "text", "text": user_message}, {"type": "image_url", "image_url": {"url": image_url}}]}],
+                "messages": [
+                    {
+                        "role": "user", 
+                        "content": [
+                            {"type": "text", "text": user_message},
+                            {"type": "image_url", "image_url": {"url": image_url}}
+                        ]
+                    }
+                ],
                 "max_tokens": 1000
             }
-            # Direct Inference URL for Vision
-            V_URL = f"https://api-inference.huggingface.co/models/{VISION_MODEL}/v1/chat/completions"
-            resp = requests.post(V_URL, headers=headers, json=payload)
-            
-            # Agar model load ho raha ho (503), toh thoda wait karke dobara try karein
-            if resp.status_code == 503:
-                time.sleep(5)
-                resp = requests.post(V_URL, headers=headers, json=payload)
-
+            resp = requests.post(ROUTER_URL, headers=headers, json=payload)
             result = resp.json()
+            
             if "choices" in result:
-                reply = result["choices"][0]["message"]["content"].strip()
+                reply = result["choices"]["message"]["content"]
             else:
-                return jsonify({"error": "Vision Model ne sahi jawab nahi diya", "api_response": result}), 500
+                return jsonify({"error": "Router Vision Fail", "api_response": result}), 500
             
-            history.append({"role": "user", "content": f"[Screenshot Sent] {user_message}"})
+            history.append({"role": "user", "content": f"[Screenshot] {user_message}"})
         else:
-            # 3. TEXT API CALL
+            # --- TEXT CALL VIA ROUTER ---
             history.append({"role": "user", "content": user_message})
             payload = {"model": TEXT_MODEL, "messages": history, "max_tokens": 1000}
-            T_URL = "https://router.huggingface.co/v1/chat/completions"
-            resp = requests.post(T_URL, headers=headers, json=payload)
-            reply = resp.json()["choices"]["message"]["content"].strip()
+            resp = requests.post(ROUTER_URL, headers=headers, json=payload)
+            reply = resp.json()["choices"][0]["message"]["content"]
 
         history.append({"role": "assistant", "content": reply})
         doc_ref.set({'messages': history})
         return jsonify({"reply": reply})
 
     except Exception as e:
-        return jsonify({"error": "System Crash!", "details": str(e)}), 500
+        return jsonify({"error": "System Crash", "details": str(e)}), 500
 
 if __name__ == '__main__':
     app.run(host='0.0.0.0', port=7860)