app.py

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()

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/',
    },
}

# 添加中间件来验证API密钥
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

# 日志记录函数
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(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)',
                    # 其他可能的参数...
                }
            }
        }
    })

# 在路由中使用API密钥验证中间件
@app.route('/ai/v1/messages', methods=['POST'])
def ai_messages():
    validation_error = validate_api_key()
    if validation_error:
        return validation_error

    raw_body = request.data.decode('utf-8')
    try:
        json_body = json.loads(raw_body)
        if json_body.get('stream') is False:
            log_request(request, 200)
            return jsonify({
                '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,
                },
            })
        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)

                    @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')
        else:
            raise ValueError('Invalid request')
    except Exception as e:
        logging.error(e)
        return jsonify({'error': str(e)}), 400

# 处理其他路由
@app.errorhandler(404)
def not_found(e):
    log_request(request, 404)
    return 'Not Found', 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}')
    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)

# eventStream util
def create_event(event, data):
    if isinstance(data, dict):
        data = json.dumps(data)
    return f'event: {event}\ndata: {data}\n\n'