import os import requests import json import threading from flask import Flask, request, jsonify, Response, stream_with_context # ================================================================================ # FLASK APP INITIALIZATION # ================================================================================ app = Flask(__name__) # ---------------------------------------------------------------------------- # CONFIGURATION & ENTERPRISE URLS # ---------------------------------------------------------------------------- # New Enterprise Gateway GAS URL GAS_URL = "https://script.google.com/macros/s/AKfycby63Fni4IEJmNJWsweBneSxbOgsYd9zBp7JskVtvKWCxaxJtjz0nnGETM6cfr1r39Ii/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" 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" } # ---------------------------------------------------------------------------- # ENTERPRISE BACKGROUND LOGGING (With Health Score & Dynamic Model Tracking) # ---------------------------------------------------------------------------- def log_usage_to_gas(email, model_name, tokens_used, status_code, api_key=""): try: payload = { "action": "log_api_usage", "email": email, "apiKey": api_key, # Passed for backward compatibility "model_name": model_name, # Dynamically tracks any model "tokens": int(tokens_used), "status_code": int(status_code) # Crucial for Health Score monitoring } # Explicitly allowing redirects to prevent 'Auth Server Offline' error requests.post(GAS_URL, json=payload, timeout=8, allow_redirects=True) except Exception: pass # Silently fail in background to keep UI blazing fast # ---------------------------------------------------------------------------- # SYSTEM PROMPT ENGINE (Identity Override) # ---------------------------------------------------------------------------- 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 # ================================================================================ # ENTERPRISE API GATEWAY ENDPOINT # ================================================================================ @app.route('/v1/chat/completions', methods=['POST']) def chat_completions(): # 1. API Key Extraction 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("Bearer ")[1].strip() # 2. Synchronous Authentication via Google Apps Script (Gateway Check) try: gas_auth_response = requests.post( GAS_URL, json={"action": "hf_verify_key", "apiKey": user_api_key}, timeout=8, allow_redirects=True ) gas_data = gas_auth_response.json() # Will block 'Suspended' or 'Deleted' keys automatically 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": f"Gateway Auth Timeout/Error.", "type": "server_error", "details": str(e)}}), 500 # 3. Request Formatting 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"]) messages = data.get("messages", []) stream = data.get("stream", False) enable_thinking = data.get("enable_thinking", True) custom_system_content = "" filtered_messages = [] 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 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) } 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" } # 4. LLM Upstream Call & Analytics Logging try: response = requests.post(INVOKE_URL, headers=headers, json=payload, stream=stream, timeout=60) # Track 4xx/5xx API Errors and send to GAS to affect Health Score if response.status_code != 200: threading.Thread(target=log_usage_to_gas, args=(user_email, requested_model, 0, response.status_code, user_api_key)).start() return jsonify({"error": {"message": f"Upstream API Error: {response.text}", "code": response.status_code}}), response.status_code # Successful Non-Streaming Response (Status 200) if not stream: nvidia_json = response.json() nvidia_json["model"] = requested_model tokens_used = nvidia_json.get("usage", {}).get("total_tokens", 0) threading.Thread(target=log_usage_to_gas, args=(user_email, requested_model, tokens_used, 200, user_api_key)).start() return jsonify(nvidia_json) # Successful Streaming Response (Status 200) def generate(): full_content_length = 0 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 if data_json.get("choices"): delta = data_json["choices"][0].get("delta", {}) full_content_length += len(delta.get("content", "")) yield "data: " + json.dumps(data_json) + "\n\n" except Exception: yield decoded + "\n\n" else: yield decoded + "\n\n" estimated_tokens = full_content_length // 4 threading.Thread(target=log_usage_to_gas, args=(user_email, requested_model, estimated_tokens, 200, user_api_key)).start() return Response(stream_with_context(generate()), mimetype='text/event-stream') except Exception as e: # Log internal server timeouts as 500 errors to GAS threading.Thread(target=log_usage_to_gas, args=(user_email, requested_model, 0, 500, user_api_key)).start() return jsonify({"error": {"message": str(e), "type": "internal_server_error"}}), 500 if __name__ == '__main__': app.run(host='0.0.0.0', port=7860)