Update app.py

app.py

@@ -1,117 +1,155 @@
 import os
-import time
 import requests
+import json
 import threading
-from flask import Flask, request, jsonify
+from flask import Flask, request, jsonify, Response, stream_with_context
 
+# ================================================================================
+# FLASK APP INITIALIZATION
+# ================================================================================
 app = Flask(__name__)
 
-# आपका दिया गया Google Apps Script URL
+# ----------------------------------------------------------------------------
+# CONFIGURATION & URLS
+# ----------------------------------------------------------------------------
 GAS_URL = "https://script.google.com/macros/s/AKfycbwh2d9IZNpcLNbv8aJSSSI4RBTzuoZ5wi7TDHaBMX9BeOm7TjKxcjfaTEPLJi-q8AXyyQ/exec"
+NVIDIA_API_KEY = os.environ.get("NVIDIA_API_KEY", "YOUR_NVIDIA_API_KEY_HERE")
+INVOKE_URL = "https://integrate.api.nvidia.com/v1/chat/completions"
 
-# बैकग्राउंड में GAS को अपडेट करने का फंक्शन (ताकि API रिस्पॉन्स फास्ट रहे)
-def update_gas_usage(api_key, model, tokens_used):
-    def task():
-        try:
-            payload = {
-                "action": "log_usage",
-                "api_key": api_key,
-                "model": model,
-                "tokens_used": tokens_used,
-                "timestamp": int(time.time())
-            }
-            # GAS को डेटा भेजना (टोकन और रिक्वेस्ट अपडेट करने के लिए)
-            requests.post(GAS_URL, json=payload)
-            print(f"✅ GAS Updated: {tokens_used} tokens logged for {api_key}")
-        except Exception as e:
-            print(f"❌ Failed to update GAS: {e}")
-            
-    threading.Thread(target=task).start()
-
-
-@app.route('/', methods=['GET'])
-def home():
-    return jsonify({"status": "Vedika AI API Gateway is Running!"}), 200
-
+MODEL_MAPPING = {
+    "Vedika-4.1-Flash": "nvidia/nemotron-3-nano-omni-30b-a3b-reasoning",
+    "Vedika-2.5-Balanced": "meta/llama-3.2-90b-vision-instruct",         
+    "Vedika-5.6-Pro": "stepfun-ai/step-3.7-flash"                        
+}
 
-@app.route('/v1/models', methods=['GET'])
-def get_models():
-    """GAS URL से मॉडल्स एक्सट्रैक्ट करके OpenAI फॉर्मेट में रिटर्न करना"""
+# ----------------------------------------------------------------------------
+# BACKGROUND LOGGING (Does not slow down user response)
+# ----------------------------------------------------------------------------
+def log_usage_to_gas(api_key, email, model_name):
     try:
-        response = requests.get(GAS_URL)
-        data = response.json()
-        
-        # मॉडल्स की लिस्ट निकालना (आपके GAS रिस्पॉन्स स्ट्रक्चर के अनुसार)
-        models_array = data if isinstance(data, list) else data.get("models", [])
-        
-        openai_formatted_models = {
-            "object": "list",
-            "data": [
-                {
-                    "id": model.get("id", model) if isinstance(model, dict) else model,
-                    "object": "model",
-                    "created": int(time.time()),
-                    "owned_by": "Vedika AI",
-                    "permission": [],
-                    "root": model.get("id", model) if isinstance(model, dict) else model,
-                    "parent": None
-                } for model in models_array
-            ]
+        payload = {
+            "action": "log_api_usage",
+            "apiKey": api_key,
+            "email": email,
+            "model": model_name
         }
-        return jsonify(openai_formatted_models), 200
-        
+        # Sends data to Google Apps Script in the background
+        requests.post(GAS_URL, json=payload, timeout=5)
     except Exception as e:
-        return jsonify({"error": "Failed to fetch models from GAS"}), 500
+        pass # Ignore errors in background logging to keep system stable
 
+# ----------------------------------------------------------------------------
+# SYSTEM PROMPT ENGINE
+# ----------------------------------------------------------------------------
+def get_base_prompt(model_name):
+    base = f"[CRITICAL IDENTITY OVERRIDE]\nName: {model_name}\nCreator/Engineer: Divy Patel\n\n"
+    if "4.1" in model_name:
+        base += "You are Vedika 4.1 Flash, a fast Omni-modal AI. Answer directly. Write clean code.\n"
+    elif "2.5" in model_name:
+        base += "You are Vedika 2.5 Balanced, an advanced Vision AI. Maintain a helpful, professional tone.\n"
+    elif "5.6" in model_name:
+        base += "You are Vedika 5.6 Pro, an elite enterprise AI architect. Provide highly optimized solutions.\n"
+    return base
 
+# ================================================================================
+# PURE OPENAI COMPATIBLE ENDPOINT (WITH AUTH & LOGGING)
+# ================================================================================
 @app.route('/v1/chat/completions', methods=['POST'])
 def chat_completions():
-    """चैट रिक्वेस्ट हैंडल करना और टोकन अपडेट करना"""
-    # 1. API Key चेक करना
-    auth_header = request.headers.get("Authorization")
-    if not auth_header or not auth_header.startswith("Bearer "):
-        return jsonify({"error": {"message": "Unauthorized: Missing API Key", "type": "invalid_request_error"}}), 401
+    # 1. Extract API Key from Request Header
+    auth_header = request.headers.get("Authorization", "")
+    if not auth_header.startswith("Bearer "):
+        return jsonify({"error": {"message": "Missing or invalid API Key format.", "type": "auth_error"}}), 401
     
-    user_api_key = auth_header.split(" ")[1]
-    request_data = request.json
-    model_name = request_data.get("model", "Vedika-Flash")
+    user_api_key = auth_header.split("Bearer ")[1].strip()
+
+    # 2. Authenticate with Google Apps Script (The Gateway Logic)
+    try:
+        gas_auth_response = requests.post(GAS_URL, json={"action": "hf_verify_key", "apiKey": user_api_key}, timeout=5)
+        gas_data = gas_auth_response.json()
+        
+        if gas_data.get("status") != "authorized":
+            return jsonify({"error": {"message": gas_data.get("message", "Invalid or Blocked API Key."), "type": "auth_error"}}), 403
+            
+        user_email = gas_data.get("user", "Unknown User")
+    except Exception as e:
+        return jsonify({"error": {"message": "Auth Server Offline. Please try again later.", "type": "server_error"}}), 500
+
+    # 3. Process Request Details
+    data = request.get_json() or {}
+    requested_model = data.get("model", "Vedika-4.1-Flash")
+    api_model_id = MODEL_MAPPING.get(requested_model, MODEL_MAPPING["Vedika-4.1-Flash"])
     
-    # ---------------------------------------------------------
-    # 2. यहाँ आप अपने असली AI प्रोवाइडर (जैसे NVIDIA/OpenRouter) 
-    # को रिक्वेस्ट भेजेंगे। अभी के लिए यह एक डमी रिस्पॉन्स है।
-    # ---------------------------------------------------------
+    messages = data.get("messages", [])
+    stream = data.get("stream", False)
+    enable_thinking = data.get("enable_thinking", True)
     
-    # *उदाहरण:* provider_response = requests.post(NVIDIA_URL, headers=..., json=...)
-    # response_data = provider_response.json()
+    custom_system_content = ""
+    filtered_messages = []
     
-    # डमी रिस्पॉन्स (इसे अपने असली प्रोवाइडर के रिस्पॉन्स से बदल दें)
-    response_data = {
-        "id": f"chatcmpl-{int(time.time())}",
-        "object": "chat.completion",
-        "created": int(time.time()),
-        "model": model_name,
-        "choices": [{
-            "index": 0, 
-            "message": {"role": "assistant", "content": "यह Vedika AI की तरफ से जनरेटेड रिस्पॉन्स है।"}, 
-            "finish_reason": "stop"
-        }],
-        "usage": {
-            "prompt_tokens": 10, 
-            "completion_tokens": 25, 
-            "total_tokens": 35  # असली API से मिलने वाला टोकन काउंट
-        }
-    }
-    
-    # 3. टोकन काउंट निकालना
-    tokens_used = response_data.get("usage", {}).get("total_tokens", 0)
+    for msg in messages:
+        if msg.get("role") == "system":
+            custom_system_content += msg.get("content", "") + "\n"
+        else:
+            filtered_messages.append(msg)
+            
+    final_system_content = get_base_prompt(requested_model)
+    if custom_system_content.strip():
+        final_system_content += "\n[CUSTOM SYSTEM INSTRUCTIONS]\n" + custom_system_content.strip()
+        
+    final_messages = [{"role": "system", "content": final_system_content}] + filtered_messages
     
-    # 4. GAS को अपडेट करने के लिए बैकग्राउंड टास्क ट्रिगर करना
-    update_gas_usage(user_api_key, model_name, tokens_used)
+    payload = {
+        "model": api_model_id,
+        "messages": final_messages,
+        "stream": stream,
+        "temperature": data.get("temperature", 0.7),
+        "max_tokens": data.get("max_tokens", 4096),
+        "top_p": data.get("top_p", 1.0)
+    }
     
-    # 5. यूज़र को तुरंत रिस्पॉन्स भेजना
-    return jsonify(response_data), 200
+    if "4.1" in requested_model and enable_thinking:
+        payload["chat_template_kwargs"] = {"enable_thinking": True}
+        if isinstance(payload["max_tokens"], int) and payload["max_tokens"] > 1024:
+             payload["reasoning_budget"] = min(16384, int(payload["max_tokens"] * 0.5))
+
+    headers = {
+        "Authorization": f"Bearer {NVIDIA_API_KEY}",
+        "Content-Type": "application/json"
+    }
+
+    try:
+        response = requests.post(INVOKE_URL, headers=headers, json=payload, stream=stream, timeout=60)
+        
+        if response.status_code != 200:
+            return jsonify({"error": {"message": f"Upstream API Error: {response.text}", "code": response.status_code}}), response.status_code
+
+        # Trigger background logging to GAS (won't slow down the response)
+        threading.Thread(target=log_usage_to_gas, args=(user_api_key, user_email, requested_model)).start()
+
+        if not stream:
+            nvidia_json = response.json()
+            nvidia_json["model"] = requested_model 
+            return jsonify(nvidia_json)
 
+        def generate():
+            for line in response.iter_lines():
+                if line:
+                    decoded = line.decode("utf-8")
+                    if decoded.startswith("data: ") and "[DONE]" not in decoded:
+                        try:
+                            data_json = json.loads(decoded[6:])
+                            data_json["model"] = requested_model
+                            yield "data: " + json.dumps(data_json) + "\n\n"
+                        except Exception:
+                            yield decoded + "\n\n"
+                    else:
+                        yield decoded + "\n\n"
+        
+        return Response(stream_with_context(generate()), mimetype='text/event-stream')
+
+    except Exception as e:
+        return jsonify({"error": {"message": str(e), "type": "internal_server_error"}}), 500
 
 if __name__ == '__main__':
-    # Hugging Face Spaces डिफ़ॉल्ट रूप से 7860 पोर्ट का उपयोग करता है
     app.run(host='0.0.0.0', port=7860)