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 app = Flask(__name__) socketio = SocketIO(app) # 从环境变量中获取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) # eventStream util def create_event(event, data): if isinstance(data, dict): data = json.dumps(data) return f'event: {event}\ndata: {data}\n\n'