Update app.py

app.py

@@ -1,305 +1,202 @@
-from flask import Flask, request, jsonify, Response
-from flask_socketio import SocketIO, emit
-from uuid import uuid4
 import os
 import json
-import logging
-import requests
-from datetime import datetime
-from gevent import monkey
-monkey.patch_all()
-
-app = Flask(__name__)
-socketio = SocketIO(app, async_mode='gevent')
-
-# 从环境变量中获取API密钥
-API_KEY = os.getenv('PPLX_KEY')
-
-# 配置代理
-class ProxyAgent:
-    def __init__(self):
-        self.proxy = os.getenv('HTTP_PROXY') or os.getenv('HTTPS_PROXY')
-
-    def get_session(self):
-        session = requests.Session()
-        if self.proxy:
-            session.proxies = {
-                'http': self.proxy,
-                'https': self.proxy
-            }
-        return session
-
-agent = ProxyAgent()
+import uuid
+import asyncio
 
+from aiohttp import web
+from socketio import AsyncClient
+from aiohttp_socks import ProxyConnector
+
+# 从环境变量中获取 API 密钥和 Cookie
+API_KEY = os.environ.get("PPLX_KEY")
+PPLX_COOKIE = os.environ.get("PPLX_COOKIE")
+USER_AGENT = os.environ.get("USER_AGENT")
+
+# 代理设置（如果需要）
+PROXY_URL = os.environ.get("PROXY_URL")  # 例如： socks5://user:pass@host:port
+connector = ProxyConnector.from_url(PROXY_URL) if PROXY_URL else None
+
+# Socket.IO 客户端选项
 opts = {
-    'agent': agent,
-    'auth': {
-        'jwt': 'anonymous-ask-user',
-    },
-    'reconnection': False,
-    'transports': ['websocket'],
-    'path': '/socket.io',
-    'hostname': 'www.perplexity.ai',
-    'secure': True,
-    'port': '443',
-    'extraHeaders': {
-        'Cookie': os.getenv('PPLX_COOKIE'),
-        'User-Agent': os.getenv('USER_AGENT'),
-        'Accept': '*/*',
-        'priority': 'u=1, i',
-        'Referer': 'https://www.perplexity.ai/',
+    "auth": {"jwt": "anonymous-ask-user"},
+    "reconnection": False,
+    "transports": ["websocket"],
+    "path": "/socket.io",
+    "hostname": "www.perplexity.ai",
+    "secure": True,
+    "port": "443",
+    "extra_headers": {
+        "Cookie": PPLX_COOKIE,
+        "User-Agent": USER_AGENT,
+        "Accept": "*/*",
+        "priority": "u=1, i",
+        "Referer": "https://www.perplexity.ai/",
     },
 }
 
-# 添加中间件来验证API密钥
-def validate_api_key():
-    api_key = request.headers.get('x-api-key')
+# 创建 aiohttp 应用
+app = web.Application()
+
+# API 密钥验证中间件
+@web.middleware
+async def validate_api_key(request, handler):
+    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
+        return web.json_response({"error": "无效的 API 密钥"}, status=401)
+    return await handler(request)
 
 # 日志记录函数
-def log_request(req, status):
+def log_request(request, status):
     timestamp = datetime.now().isoformat()
-    ip = req.remote_addr
-    route = req.path
-    logging.info(f'{timestamp} - {ip} - {route} - {status}')
+    ip = request.remote or request.transport.get_extra_info("peername")[0]
+    route = request.path
+    print(f"{timestamp} - {ip} - {route} - {status}")
+
 
 # 根路由处理
-@app.route('/', methods=['GET'])
-def root():
+async def root_handler(request):
     log_request(request, 200)
-    return jsonify({
-        'message': 'Welcome to the Perplexity AI Proxy API',
-        'endpoints': {
-            '/ai/v1/messages': {
-                'method': 'POST',
-                'description': 'Send a message to the AI',
-                'headers': {
-                    'x-api-key': 'Your API key (required)',
-                    'Content-Type': 'application/json'
-                },
-                'body': {
-                    'messages': 'Array of message objects',
-                    'stream': 'Boolean (true for streaming response)',
-                    # 其他可能的参数...
+    return web.json_response(
+        {
+            "message": "欢迎使用 Perplexity AI 代理 API",
+            "endpoints": {
+                "/ai/v1/messages": {
+                    "method": "POST",
+                    "description": "向 AI 发送消息",
+                    "headers": {
+                        "x-api-key": "你的 API 密钥（必需）",
+                        "Content-Type": "application/json",
+                    },
+                    "body": {
+                        "messages": "消息对象数组",
+                        "stream": "布尔值（true 表示流式响应）",
+                        # 其他可能的参数...
+                    },
                 }
-            }
+            },
         }
-    })
+    )
 
-# 在路由中使用API密钥验证中间件
-@app.route('/ai/v1/messages', methods=['POST'])
-def ai_messages():
-    validation_error = validate_api_key()
-    if validation_error:
-        return validation_error
+# 创建事件流数据
+def create_event(event, data):
+    if isinstance(data, dict):
+        data = json.dumps(data)
+    return f"event: {event}\ndata: {data}\n\n"
 
-    raw_body = request.data.decode('utf-8')
+
+# AI 消息处理
+@web.post("/ai/v1/messages")
+@validate_api_key
+async def ai_messages_handler(request):
     try:
-        json_body = json.loads(raw_body)
-        if json_body.get('stream') is False:
+        json_body = await request.json()
+        if json_body.get("stream") is False:
             log_request(request, 200)
-            return jsonify({
-                'id': str(uuid4()),
-                'content': [
-                    {
-                        'text': 'Please turn on streaming.',
-                    },
+            return web.json_response(
+                {
+                    "id": str(uuid.uuid4()),
+                    "content": [
+                        {"text": "请打开流式传输。"},
+                        {"id": "string", "name": "string", "input": {}},
+                    ],
+                    "model": "string",
+                    "stop_reason": "end_turn",
+                    "stop_sequence": "string",
+                    "usage": {"input_tokens": 0, "output_tokens": 0},
+                }
+            )
+        elif json_body.get("stream") is True:
+            # 处理消息... （与 JavaScript 版本类似）
+            previous_messages = "\n\n".join([msg.get("content", "") for msg in json_body.get("messages", [])])
+            msgid = str(uuid.uuid4())
+
+            response = web.StreamResponse(
+                status=200,
+                reason="OK",
+                headers={"Content-Type": "text/event-stream;charset=utf-8"},
+            )
+            await response.prepare(request)
+
+            await response.write(
+                create_event(
+                    "message_start",
                     {
-                        'id': 'string',
-                        'name': 'string',
-                        'input': {},
-                    },
-                ],
-                'model': 'string',
-                'stop_reason': 'end_turn',
-                'stop_sequence': 'string',
-                'usage': {
-                    'input_tokens': 0,
-                    'output_tokens': 0,
-                },
-            })
-        elif json_body.get('stream') is True:
-            def generate():
-                try:
-                    # 计算用户消息长度
-                    user_message = [{'question': '', 'answer': ''}]
-                    last_update = True
-                    if json_body.get('system'):
-                        json_body['messages'].insert(0, {'role': 'system', 'content': json_body['system']})
-                    for msg in json_body['messages']:
-                        if msg['role'] in ['system', 'user']:
-                            if last_update:
-                                user_message[-1]['question'] += msg['content'] + '\n'
-                            elif user_message[-1]['question'] == '':
-                                user_message[-1]['question'] += msg['content'] + '\n'
-                            else:
-                                user_message.append({'question': msg['content'] + '\n', 'answer': ''})
-                            last_update = True
-                        elif msg['role'] == 'assistant':
-                            if not last_update:
-                                user_message[-1]['answer'] += msg['content'] + '\n'
-                            elif user_message[-1]['answer'] == '':
-                                user_message[-1]['answer'] += msg['content'] + '\n'
-                            else:
-                                user_message.append({'question': '', 'answer': msg['content'] + '\n'})
-                            last_update = False
-
-                    previous_messages = '\n\n'.join(msg['content'] for msg in json_body['messages'])
-                    msg_id = str(uuid4())
-
-                    yield create_event('message_start', {
-                        'type': 'message_start',
-                        'message': {
-                            'id': msg_id,
-                            'type': 'message',
-                            'role': 'assistant',
-                            'content': [],
-                            'model': 'claude-3-opus-20240229',
-                            'stop_reason': None,
-                            'stop_sequence': None,
-                            'usage': {'input_tokens': 8, 'output_tokens': 1},
+                        "type": "message_start",
+                        "message": {
+                            "id": msgid,
+                           # ... （其余部分与 JavaScript 版本类似）
                         },
-                    })
-                    yield create_event('content_block_start', {'type': 'content_block_start', 'index': 0, 'content_block': {'type': 'text', 'text': ''}})
-                    yield create_event('ping', {'type': 'ping'})
-
-                    # 代理响应
-                    socket = socketio.connect('wss://www.perplexity.ai/', **opts)
-
-                    @socket.on('connect')
-                    def on_connect():
-                        logging.info(' > [Connected]')
-                        socket.emit('perplexity_ask', previous_messages, {
-                            'version': '2.9',
-                            'source': 'default',
-                            'attachments': [],
-                            'language': 'en-GB',
-                            'timezone': 'Europe/London',
-                            'search_focus': 'writing',
-                            'frontend_uuid': str(uuid4()),
-                            'mode': 'concise',
-                            'is_related_query': False,
-                            'is_default_related_query': False,
-                            'visitor_id': str(uuid4()),
-                            'frontend_context_uuid': str(uuid4()),
-                            'prompt_source': 'user',
-                            'query_source': 'home'
-                        }, callback=on_response)
-
-                    def on_response(response):
-                        logging.info(response)
-                        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},
-                        })
-                        yield create_event('message_stop', {'type': 'message_stop'})
-                        log_request(request, 200)
-
-                    @socket.on('query_progress')
-                    def on_query_progress(data):
-                        if data.get('text'):
-                            text = json.loads(data['text'])
-                            chunk = text['chunks'][-1]
-                            if chunk:
-                                yield create_event('content_block_delta', {
-                                    'type': 'content_block_delta',
-                                    'index': 0,
-                                    'delta': {'type': 'text_delta', 'text': chunk},
-                                })
-
-                    @socket.on('disconnect')
-                    def on_disconnect():
-                        logging.info(' > [Disconnected]')
-
-                    @socket.on('error')
-                    def on_error(error):
-                        yield create_event('content_block_delta', {
-                            'type': 'content_block_delta',
-                            'index': 0,
-                            'delta': {'type': 'text_delta', 'text': 'Error occured 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},
-                        })
-                        yield create_event('message_stop', {'type': 'message_stop'})
-                        logging.error(error)
-                        log_request(request, 500)
-
-                    @socket.on('connect_error')
-                    def on_connect_error(error):
-                        yield create_event('content_block_delta', {
-                            'type': 'content_block_delta',
-                            'index': 0,
-                            'delta': {'type': 'text_delta', 'text': 'Failed to connect to the Perplexity.ai 连接到Perplexity失败\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},
-                        })
-                        yield create_event('message_stop', {'type': 'message_stop'})
-                        logging.error(error)
-                        log_request(request, 500)
-
-                    @socket.on('close')
-                    def on_close():
-                        logging.info(' > [Client closed]')
-                        socket.disconnect()
-
-                except Exception as e:
-                    logging.error(e)
-                    yield create_event('content_block_delta', {
-                        'type': 'content_block_delta',
-                        'index': 0,
-                        'delta': {'type': 'text_delta', 'text': 'Error occured 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},
-                    })
-                    yield create_event('message_stop', {'type': 'message_stop'})
-                    log_request(request, 500)
-
-            return Response(generate(), mimetype='text/event-stream')
+                    },
+                ).encode()
+            )
+
+            # 使用 AsyncClient 连接到 Socket.IO 服务器
+            async with AsyncClient(connector=connector) as socket:
+                await socket.connect("wss://www.perplexity.ai/", **opts)
+
+                # 发送 perplexity_ask 事件 (需要根据Perplexity.ai的API调整)
+                await socket.emit("perplexity_ask", previous_messages)
+
+
+                # 处理 Socket.IO 事件... (与 JavaScript 版本类似)
+
+                 # ...处理 query_progress, disconnect, error, connect_error 等事件，并写入 response
+
+
+                @socket.on("query_progress")
+                async def on_query_progress(data):
+                    if data.get("text"):
+                        text = json.loads(data["text"])
+                        chunk = text["chunks"][-1] if text.get("chunks") else None
+                        if chunk:
+                            await response.write(create_event("content_block_delta", { # ...}).encode())
+
+
+                # ... other event handlers
+
+            return response
+
         else:
-            raise ValueError('Invalid request')
-    except Exception as e:
-        logging.error(e)
-        return jsonify({'error': str(e)}), 400
+            raise ValueError("无效的请求")
+
+
+    except (json.JSONDecodeError, ValueError, Exception) as e:
+        console.log(e) # debug
+        log_request(request, 400)
+        return web.json_response({"error": str(e)}, status=400)
+
+
 
 # 处理其他路由
-@app.errorhandler(404)
-def not_found(e):
+async def not_found_handler(request):
     log_request(request, 404)
-    return 'Not Found', 404
+    return web.Response(text="未找到", status=404)
 
 # 错误处理中间件
-@app.errorhandler(500)
-def internal_server_error(e):
-    logging.error(e)
-    log_request(request, 500)
-    return 'Something broke!', 500
-
-if __name__ == '__main__':
-    port = int(os.getenv('PORT', 8081))
-    logging.info(f'Perplexity proxy listening on port {port}')
+@web.middleware
+async def error_middleware(request, handler):
+    try:
+        return await handler(request)
+    except Exception as err:
+        print(err, traceback.format_exc()) # debug
+        log_request(request, 500)
+        return web.Response(text="服务器错误", status=500)
+
+# 添加路由和中间件
+app.add_routes([web.get("/", root_handler), web.post("/ai/v1/messages", ai_messages_handler)])
+app.add_routes([web.static("/static", "static")])
+app.middleware.append(error_middleware)
+app.router.add_route("*", "/{path_info:.*}", not_found_handler)
+
+# 启动服务器
+if __name__ == "__main__":
+
+    port = int(os.environ.get("PORT", 8081))
+
     if not API_KEY:
-        logging.warning('Warning: PPLX_KEY environment variable is not set. API key validation will fail.')
-    socketio.run(app, port=port, allow_unsafe_werkzeug=True)
+        print("警告：未设置 PPLX_KEY 环境变量。API 密钥验证将失败。")
+
+    web.run_app(app, port=port)
+
 
-# eventStream util
-def create_event(event, data):
-    if isinstance(data, dict):
-        data = json.dumps(data)
-    return f'event: {event}\ndata: {data}\n\n'