Update app.py

app.py

@@ -5,95 +5,99 @@ import json
 
 app = Flask(__name__)
 
-# Target API base URL from environment variable
-TARGET_API = os.getenv("TARGET_API", "https://huggingface.co")
+# --- Configuration from Environment Variables ---
 
-# Path mappings from environment variable
-# Expected format: {"path1": "mapped_path1", "path2": "mapped_path2"}
+# 1. Target API base URL
+TARGET_API = os.getenv("TARGET_API", "https://api-inference.huggingface.co")
+
+# 2. The REAL secret key for the target API. This is kept on the server.
+REAL_AUTH_KEY = os.getenv("REAL_AUTH_KEY")
+
+# 3. The access key for this proxy. This is what you share with your friends.
+PROXY_ACCESS_KEY = os.getenv("PROXY_ACCESS_KEY")
+
+
+# 4. Path mappings from environment variable
+# Example: '{"/v1/chat/completions": "/models/mistralai/Mixtral-8x7B-Instruct-v0.1"}'
 def get_path_mappings():
-    mappings_str = os.getenv("PATH_MAPPINGS", '{"/": "/"}')
+    mappings_str = os.getenv("PATH_MAPPINGS", '{}') # Default to empty dict
     try:
         return json.loads(mappings_str)
     except json.JSONDecodeError:
-        # Fallback to default mappings if JSON is invalid
-        return {
-            "/": "/",
-        }
+        print("Warning: Invalid JSON in PATH_MAPPINGS. Using empty mappings.")
+        return {}
 
 PATH_MAPPINGS = get_path_mappings()
 
 
 @app.route('/<path:path>', methods=['GET', 'POST', 'PUT', 'DELETE', 'PATCH'])
 def proxy(path):
-    # Construct the full path
+    # --- 1. Authentication Check ---
+    # Ensure the server is configured with the necessary keys
+    if not REAL_AUTH_KEY or not PROXY_ACCESS_KEY:
+        error_msg = {"error": "Authentication is not configured on the proxy server."}
+        return Response(json.dumps(error_msg), status=500, mimetype='application/json')
+
+    # Get the authorization header from the user's request
+    auth_header = request.headers.get('Authorization')
+    expected_auth_header = f"Bearer {PROXY_ACCESS_KEY}"
+
+    # Validate the proxy access key
+    if auth_header != expected_auth_header:
+        error_msg = {"error": "Invalid or missing proxy access key."}
+        return Response(json.dumps(error_msg), status=401, mimetype='application/json')
+
+    # --- 2. Path and URL Construction ---
     full_path = f"/{path}"
 
-    # Apply path mapping if matches
+    # Apply path mapping if a match is found
+    # This allows you to map a generic path like /v1/chat to a specific model path
     for original_path, new_path in PATH_MAPPINGS.items():
-        if full_path.startswith(original_path):
-            full_path = full_path.replace(original_path, new_path, 1)
+        if full_path == original_path:
+            full_path = new_path
             break
 
-    # Construct target URL
     target_url = f"{TARGET_API}{full_path}"
 
-    # Forward the request to the target API
-    headers = {key: value for key, value in request.headers if key != 'Host'}
+    # --- 3. Header Manipulation ---
+    # Copy headers from the incoming request, but remove 'Host' and the user's 'Authorization'
+    headers = {key: value for key, value in request.headers if key.lower() not in ['host', 'authorization']}
+    
+    # Add the REAL authentication key for the target API
+    headers['Authorization'] = f"Bearer {REAL_AUTH_KEY}"
 
-    # Handle streaming response
-    if request.method == 'POST':
-        response = requests.post(
-            target_url,
-            headers=headers,
-            json=request.get_json(silent=True),
-            params=request.args,
-            stream=True
-        )
-    elif request.method == 'GET':
-        response = requests.get(
-            target_url,
-            headers=headers,
-            params=request.args,
-            stream=True
-        )
-    elif request.method == 'PUT':
-        response = requests.put(
-            target_url,
-            headers=headers,
-            json=request.get_json(silent=True),
-            params=request.args,
-            stream=True
-        )
-    elif request.method == 'DELETE':
-        response = requests.delete(
-            target_url,
-            headers=headers,
-            params=request.args,
-            stream=True
-        )
-    elif request.method == 'PATCH':
-        response = requests.patch(
-            target_url,
+    # --- 4. Forward the Request (Refactored for all methods) ---
+    try:
+        # Use request.get_data() to handle any kind of request body (JSON, form data, etc.)
+        response = requests.request(
+            method=request.method,
+            url=target_url,
             headers=headers,
-            json=request.get_json(silent=True),
             params=request.args,
+            data=request.get_data(),
             stream=True
         )
+    except requests.exceptions.RequestException as e:
+        error_msg = {"error": f"Failed to connect to target service: {e}"}
+        return Response(json.dumps(error_msg), status=502, mimetype='application/json')
+
 
-    # Create a response with the same status code, headers, and streaming content
+    # --- 5. Stream the Response Back to the Client ---
     def generate():
         for chunk in response.iter_content(chunk_size=8192):
             yield chunk
 
-    # Create flask response
+    # Create a Flask response object
     proxy_response = Response(
         stream_with_context(generate()),
         status=response.status_code
     )
 
-    # Forward response headers
+    # Copy headers from the target's response to our proxy response
+    # Exclude certain headers that are handled by the WSGI server
+    excluded_headers = ['content-encoding', 'content-length', 'transfer-encoding', 'connection']
     for key, value in response.headers.items():
-        if key.lower() not in ('content-length', 'transfer-encoding', 'connection'):
+        if key.lower() not in excluded_headers:
             proxy_response.headers[key] = value
 
     return proxy_response
@@ -101,8 +105,9 @@ def proxy(path):
 
 @app.route('/', methods=['GET'])
 def index():
-    return "service running."
+    return "Proxy service is running."
 
 
 if __name__ == '__main__':
-    app.run(host='0.0.0.0', port=7860, debug=False)
+    # It's recommended to run this with a production-grade WSGI server like Gunicorn or uWSGI
+    app.run(host='0.0.0.0', port=7860, debug=False)