Update app.py

app.py

@@ -1,245 +1,303 @@
+from flask import Flask, request, jsonify, Response
+from flask_socketio import SocketIO, emit
+from uuid import uuid4
 import os
 import json
-import uuid
-from datetime import datetime
-from flask import Flask, request, Response, jsonify
-import socketio
-import requests
 import logging
-from threading import Event
-from contextlib import contextmanager
+import requests
+from datetime import datetime
 
 app = Flask(__name__)
-
-# 从环境变量中获取日志级别,默认为INFO
-log_level = os.environ.get('LOG_LEVEL', 'INFO').upper()
-logging.basicConfig(level=getattr(logging, log_level))
+socketio = SocketIO(app)
 
 # 从环境变量中获取API密钥
-API_KEY = os.environ.get('PPLX_KEY')
-
-# 代理设置
-proxy_url = os.environ.get('PROXY_URL')
-
-# 设置代理
-if proxy_url:
-    proxies = {
-        'http': proxy_url,
-        'https': proxy_url
-    }
-    transport = requests.Session()
-    transport.proxies.update(proxies)
-else:
-    transport = None
-
-sio = socketio.Client(http_session=transport, logger=True, engineio_logger=True)
-
-# 连接选项
-connect_opts = {
-    'transports': ['websocket', 'polling'],  # 允许回退到轮询
-}
+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
 
-# 其他选项
-sio_opts = {
+agent = ProxyAgent()
+
+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.environ.get('PPLX_COOKIE'),
-        'User-Agent': os.environ.get('USER_AGENT'),
+        'Cookie': os.getenv('PPLX_COOKIE'),
+        'User-Agent': os.getenv('USER_AGENT'),
         'Accept': '*/*',
         'priority': 'u=1, i',
         'Referer': 'https://www.perplexity.ai/',
-    }
+    },
 }
 
-def log_request(timestamp, ip, route, status, request_id):
-    logging.info(f"{timestamp} - {ip} - {route} - {status} - Request ID: {request_id}")
-
+# 添加中间件来验证API密钥
 def validate_api_key():
     api_key = request.headers.get('x-api-key')
     if api_key != API_KEY:
-        log_request(datetime.now().isoformat(), request.remote_addr, request.path, 401, str(uuid.uuid4()))
-        return jsonify({"error": "Invalid API key"}), 401
+        log_request(request, 401)
+        return jsonify({'error': 'Invalid API key'}), 401
     return None
 
-@app.route('/')
+# 日志记录函数
+def log_request(req, status):
+    timestamp = datetime.now().isoformat()
+    ip = req.remote_addr
+    route = req.path
+    logging.info(f'{timestamp} - {ip} - {route} - {status}')
+
+# 根路由处理
+@app.route('/', methods=['GET'])
 def root():
-    log_request(datetime.now().isoformat(), request.remote_addr, request.path, 200, str(uuid.uuid4()))
+    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"
+        '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)",
+                'body': {
+                    'messages': 'Array of message objects',
+                    'stream': 'Boolean (true for streaming response)',
+                    # 其他可能的参数...
                 }
             }
         }
     })
 
+# 在路由中使用API密钥验证中间件
 @app.route('/ai/v1/messages', methods=['POST'])
-def messages():
-    auth_error = validate_api_key()
-    if auth_error:
-        return auth_error
+def ai_messages():
+    validation_error = validate_api_key()
+    if validation_error:
+        return validation_error
 
+    raw_body = request.data.decode('utf-8')
     try:
-        json_body = request.json
-        if not json_body.get('stream', False):
-            log_request(datetime.now().isoformat(), request.remote_addr, request.path, 200, str(uuid.uuid4()))
+        json_body = json.loads(raw_body)
+        if json_body.get('stream') is False:
+            log_request(request, 200)
             return jsonify({
-                "id": str(uuid.uuid4()),
-                "content": [
-                    {"text": "Please turn on streaming."},
-                    {"id": "string", "name": "string", "input": {}}
+                'id': str(uuid4()),
+                'content': [
+                    {
+                        'text': 'Please turn on streaming.',
+                    },
+                    {
+                        'id': 'string',
+                        'name': 'string',
+                        'input': {},
+                    },
                 ],
-                "model": "string",
-                "stop_reason": "end_turn",
-                "stop_sequence": "string",
-                "usage": {"input_tokens": 0, "output_tokens": 0}
-            })
-        
-        # 捕获请求信息
-        request_timestamp = datetime.now().isoformat()
-        request_ip = request.remote_addr
-        request_path = request.path
-        request_id = str(uuid.uuid4())
-        
-        def generate():
-            log_request(request_timestamp, request_ip, request_path, 200, request_id)
-            
-            previous_messages = "\n\n".join([msg['content'] for msg in json_body['messages']])
-            msg_id = str(uuid.uuid4())
-            response_event = Event()
-            response_text = []
-
-            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},
+                'model': 'string',
+                'stop_reason': 'end_turn',
+                'stop_sequence': 'string',
+                'usage': {
+                    'input_tokens': 0,
+                    'output_tokens': 0,
                 },
             })
-            yield create_event("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}})
-            yield create_event("ping", {"type": "ping"})
-
-            @sio.event
-            def connect():
-                logging.info("Connected to Perplexity AI")
-                sio.emit('perplexity_ask', {
-                    "text": previous_messages,
-                    "version": "2.9",
-                    "source": "default",
-                    "attachments": [],
-                    "language": "en-GB",
-                    "timezone": "Europe/London",
-                    "search_focus": "writing",
-                    "frontend_uuid": str(uuid.uuid4()),
-                    "mode": "concise",
-                    "is_related_query": False,
-                    "is_default_related_query": False,
-                    "visitor_id": str(uuid.uuid4()),
-                    "frontend_context_uuid": str(uuid.uuid4()),
-                    "prompt_source": "user",
-                    "query_source": "home"
-                })
-
-            @sio.event
-            def query_progress(data):
-                if 'text' in data:
-                    text = json.loads(data['text'])
-                    chunk = text['chunks'][-1] if text['chunks'] else None
-                    if chunk:
-                        response_text.append(chunk)
-                        yield create_event("content_block_delta", {
-                            "type": "content_block_delta",
-                            "index": 0,
-                            "delta": {"type": "text_delta", "text": chunk},
+        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},
+                        },
+                    })
+                    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)
 
-            @sio.event
-            def query_complete(data):
-                response_event.set()
-
-            @sio.event
-            def disconnect():
-                logging.info("Disconnected from Perplexity AI")
-                response_event.set()
-
-            @sio.event
-            def connect_error(data):
-                logging.error(f"Connection error: {data}")
-                yield create_event("content_block_delta", {
-                    "type": "content_block_delta",
-                    "index": 0,
-                    "delta": {"type": "text_delta", "text": f"Error connecting to Perplexity AI: {data}"},
-                })
-                response_event.set()
-
-            @contextmanager
-            def socket_connection():
-                try:
-                    sio.connect('wss://www.perplexity.ai/', **connect_opts, headers=sio_opts['extraHeaders'])
-                    yield
-                finally:
-                    if sio.connected:
-                        sio.disconnect()
-
-            try:
-                with socket_connection():
-                    response_event.wait(timeout=30)  # 等待响应，最多30秒
-            except Exception as e:
-                logging.error(f"Error during socket connection: {str(e)}")
-                yield create_event("content_block_delta", {
-                    "type": "content_block_delta",
-                    "index": 0,
-                    "delta": {"type": "text_delta", "text": f"Error during socket connection: {str(e)}"},
-                })
-
-            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": len(''.join(response_text))},
-            })
-            yield create_event("message_stop", {"type": "message_stop"})
+                    @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},
+                                })
 
-        return Response(generate(), content_type='text/event-stream')
+                    @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')
+        else:
+            raise ValueError('Invalid request')
     except Exception as e:
-        logging.error(f"Request error: {str(e)}")
-        log_request(datetime.now().isoformat(), request.remote_addr, request.path, 400, str(uuid.uuid4()))
-        return jsonify({"error": str(e)}), 400
+        logging.error(e)
+        return jsonify({'error': str(e)}), 400
 
+# 处理其他路由
 @app.errorhandler(404)
-def not_found(error):
-    log_request(datetime.now().isoformat(), request.remote_addr, request.path, 404, str(uuid.uuid4()))
-    return "Not Found", 404
+def not_found(e):
+    log_request(request, 404)
+    return 'Not Found', 404
 
-@app.errorhandler(Exception)
-def handle_exception(e):
-    logging.error(f"Unhandled exception: {str(e)}")
-    log_request(datetime.now().isoformat(), request.remote_addr, request.path, 500, str(uuid.uuid4()))
-    return jsonify({"error": "Internal server error"}), 500
+# 错误处理中间件
+@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}')
+    if not API_KEY:
+        logging.warning('Warning: PPLX_KEY environment variable is not set. API key validation will fail.')
+    socketio.run(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"
-
-if __name__ == '__main__':
-    port = int(os.environ.get('PORT', 8081))
-    logging.info(f"Perplexity proxy listening on port {port}")
-    if not API_KEY:
-        logging.warning("Warning: PPLX_KEY environment variable is not set. API key validation will fail.")
-    app.run(host='0.0.0.0', port=port)
+    return f'event: {event}\ndata: {data}\n\n'