Update app.py

app.py

@@ -1,37 +1,52 @@
 import os
 import json
 import uuid
-from flask import Flask, request, Response, jsonify, stream_with_context
-from flask_limiter import Limiter
-from flask_limiter.util import get_remote_address
+from flask import Flask, request, Response, jsonify
 import socketio
-import requests
-from requests.exceptions import RequestException
+from aiohttp import web
+from aiohttp_socks import ProxyConnector
 
 app = Flask(__name__)
+sio = socketio.Client()
 
-# 从环境变量中获取配置
+# 从环境变量中获取API密钥
 API_KEY = os.environ.get('PPLX_KEY')
-PPLX_COOKIE = os.environ.get('PPLX_COOKIE')
-USER_AGENT = os.environ.get('USER_AGENT')
-PROXY = os.environ.get('PROXY')
-
-# 设置限流
-limiter = Limiter(
-    get_remote_address,
-    app=app,
-    default_limits=["100 per 5 minutes"],
-    storage_uri="memory://"
-)
-
-# Socket.IO 客户端设置
-sio = socketio.Client()
 
 # 代理设置
-proxies = {'https': PROXY} if PROXY else None
+proxy_url = os.environ.get('PROXY_URL')
+connector = ProxyConnector.from_url(proxy_url) if proxy_url else None
+
+sio_opts = {
+    'auth': {
+        'jwt': 'anonymous-ask-user',
+    },
+    'reconnection': False,
+    'transports': ['websocket'],
+    'extraHeaders': {
+        'Cookie': os.environ.get('PPLX_COOKIE'),
+        'User-Agent': os.environ.get('USER_AGENT'),
+        'Accept': '*/*',
+        'priority': 'u=1, i',
+        'Referer': 'https://www.perplexity.ai/',
+    }
+}
+
+def log_request(req, status):
+    timestamp = datetime.now().isoformat()
+    ip = req.remote_addr
+    route = req.path
+    print(f"{timestamp} - {ip} - {route} - {status}")
+
+def validate_api_key():
+    api_key = request.headers.get('x-api-key')
+    if api_key != API_KEY:
+        log_request(request, 401)
+        return jsonify({"error": "Invalid API key"}), 401
+    return None
 
 @app.route('/')
-def home():
+def root():
+    log_request(request, 200)
     return jsonify({
         "message": "Welcome to the Perplexity AI Proxy API",
         "endpoints": {
@@ -50,25 +65,16 @@ def home():
         }
     })
 
-@app.route('/health')
-def health_check():
-    return jsonify({"status": "OK"}), 200
-
-def validate_api_key():
-    api_key = request.headers.get('x-api-key')
-    if api_key != API_KEY:
-        return jsonify({"error": "Invalid API key"}), 401
-
 @app.route('/ai/v1/messages', methods=['POST'])
-@limiter.limit("100 per 15 minutes")
-def ai_messages():
-    auth_result = validate_api_key()
-    if auth_result:
-        return auth_result
+def messages():
+    auth_error = validate_api_key()
+    if auth_error:
+        return auth_error
 
     try:
-        data = request.json
-        if not data.get('stream', False):
+        json_body = request.json
+        if not json_body.get('stream', False):
+            log_request(request, 200)
             return jsonify({
                 "id": str(uuid.uuid4()),
                 "content": [
@@ -80,9 +86,9 @@ def ai_messages():
                 "stop_sequence": "string",
                 "usage": {"input_tokens": 0, "output_tokens": 0}
             })
-
+        
         def generate():
-            previous_messages = "\n\n".join([msg['content'] for msg in data['messages']])
+            previous_messages = "\n\n".join([msg['content'] for msg in json_body['messages']])
             msg_id = str(uuid.uuid4())
 
             yield create_event("message_start", {
@@ -101,14 +107,8 @@ def ai_messages():
             yield create_event("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}})
             yield create_event("ping", {"type": "ping"})
 
-            try:
-                sio.connect('wss://www.perplexity.ai', transports=['websocket'], 
-                            headers={
-                                'Cookie': PPLX_COOKIE,
-                                'User-Agent': USER_AGENT
-                            },
-                            http_session=requests.Session() if PROXY else None,
-                            proxies=proxies)
+            @sio.on('connect')
+            def on_connect():
                 sio.emit('perplexity_ask', previous_messages, {
                     "version": "2.9",
                     "source": "default",
@@ -126,29 +126,42 @@ def ai_messages():
                     "query_source": "home"
                 })
 
-                @sio.on('query_progress')
-                def on_query_progress(data):
-                    if data.get('text'):
-                        text = json.loads(data['text'])
-                        chunk = text['chunks'][-1] if text['chunks'] else None
-                        if chunk:
-                            yield create_event("content_block_delta", {
-                                "type": "content_block_delta",
-                                "index": 0,
-                                "delta": {"type": "text_delta", "text": chunk},
-                            })
-
-                sio.wait()
-                
-            except Exception as e:
-                print(f"Socket error: {e}")
+            @sio.on('query_progress')
+            def on_query_progress(data):
+                if 'text' in data:
+                    text = json.loads(data['text'])
+                    chunk = text['chunks'][-1] if text['chunks'] else None
+                    if chunk:
+                        yield create_event("content_block_delta", {
+                            "type": "content_block_delta",
+                            "index": 0,
+                            "delta": {"type": "text_delta", "text": chunk},
+                        })
+
+            @sio.on('disconnect')
+            def on_disconnect():
+                print(" > [Disconnected]")
+
+            @sio.on('error')
+            def on_error(error):
+                print(f"Socket error: {error}")
                 yield create_event("content_block_delta", {
                     "type": "content_block_delta",
                     "index": 0,
-                    "delta": {"type": "text_delta", "text": "An error occurred while processing your request"},
+                    "delta": {"type": "text_delta", "text": "Error occurred while fetching output 输出时出现错误\nPlease refer to the log for more information 请查看日志以获取更多信息"},
+                })
+                yield create_event("content_block_stop", {"type": "content_block_stop", "index": 0})
+                yield create_event("message_delta", {
+                    "type": "message_delta",
+                    "delta": {"stop_reason": "end_turn", "stop_sequence": None},
+                    "usage": {"output_tokens": 12},
                 })
-            finally:
-                sio.disconnect()
+                yield create_event("message_stop", {"type": "message_stop"})
+                log_request(request, 500)
+
+            sio.connect('wss://www.perplexity.ai/', **sio_opts)
+            sio.wait()
+            sio.disconnect()
 
             yield create_event("content_block_stop", {"type": "content_block_stop", "index": 0})
             yield create_event("message_delta", {
@@ -157,26 +170,34 @@ def ai_messages():
                 "usage": {"output_tokens": 12},
             })
             yield create_event("message_stop", {"type": "message_stop"})
+            log_request(request, 200)
 
-        return Response(stream_with_context(generate()), content_type='text/event-stream')
+        return Response(generate(), content_type='text/event-stream')
 
     except Exception as e:
-        print(f"Request error: {e}")
+        print(f"Request error: {str(e)}")
+        log_request(request, 400)
         return jsonify({"error": str(e)}), 400
 
-def create_event(event, data):
-    if isinstance(data, dict):
-        data = json.dumps(data)
-    return f"event: {event}\ndata: {data}\n\n"
-
 @app.errorhandler(404)
 def not_found(error):
-    return jsonify({"error": "Not Found"}), 404
+    log_request(request, 404)
+    return "Not Found", 404
 
 @app.errorhandler(500)
 def server_error(error):
-    return jsonify({"error": "Internal Server Error"}), 500
+    print(f"Server error: {str(error)}")
+    log_request(request, 500)
+    return "Something broke!", 500
+
+def create_event(event, data):
+    if isinstance(data, dict):
+        data = json.dumps(data)
+    return f"event: {event}\ndata: {data}\n\n"
 
 if __name__ == '__main__':
     port = int(os.environ.get('PORT', 8081))
+    print(f"Perplexity proxy listening on port {port}")
+    if not API_KEY:
+        print("Warning: PPLX_KEY environment variable is not set. API key validation will fail.")
     app.run(host='0.0.0.0', port=port)