import os import json import uuid import re from datetime import datetime from flask import Flask, request, Response, jsonify import socketio import requests import logging from threading import Event app = Flask(__name__) # 从环境变量中获取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=False, engineio_logger=False) # 连接选项 connect_opts = { 'transports': ['websocket', 'polling'], } # 其他选项 sio_opts = { 'extraHeaders': { 'Cookie': os.environ.get('PPLX_COOKIE'), 'User-Agent': os.environ.get('USER_AGENT'), 'Accept': '*/*', 'priority': 'u=1, i', 'Referer': 'https://www.perplexity.ai/', } } class CustomFormatter(logging.Formatter): def format(self, record): log_data = { "timestamp": self.formatTime(record, self.datefmt), "level": record.levelname, "message": self.remove_ansi_escape(record.getMessage()), } if hasattr(record, 'event_type'): log_data['event_type'] = record.event_type if hasattr(record, 'data'): log_data['data'] = record.data return json.dumps(log_data, ensure_ascii=False, indent=2) def remove_ansi_escape(self, text): ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])') return ansi_escape.sub('', text) def setup_logging(): logger = logging.getLogger() logger.setLevel(logging.INFO) handler = logging.StreamHandler() handler.setFormatter(CustomFormatter()) logger.addHandler(handler) logger = logging.getLogger(__name__) def log_request(ip, route, status): timestamp = datetime.now().isoformat() logger.info(f"Request received", extra={ 'event_type': 'request', 'data': { 'timestamp': timestamp, 'ip': ip, 'route': route, 'status': status } }) def validate_api_key(): api_key = request.headers.get('x-api-key') if api_key != API_KEY: log_request(request.remote_addr, request.path, 401) return jsonify({"error": "Invalid API key"}), 401 return None def normalize_content(content): if isinstance(content, str): return content elif isinstance(content, dict): return json.dumps(content, ensure_ascii=False) elif isinstance(content, list): return " ".join([normalize_content(item) for item in content]) else: return "" def calculate_tokens(text): if re.search(r'[^\x00-\x7F]', text): return len(text) else: tokens = text.split() return len(tokens) def create_event(event, data): if isinstance(data, dict): data = json.dumps(data, ensure_ascii=False) return f"event: {event}\ndata: {data}\n\n" @app.route('/') def root(): log_request(request.remote_addr, request.path, 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)", "model": "Model to be used (optional, defaults to claude-3-opus-20240229)" } } } }) @app.route('/ai/v1/messages', methods=['POST']) def messages(): auth_error = validate_api_key() if auth_error: return auth_error try: json_body = request.json model = json_body.get('model', 'claude-3-opus-20240229') stream = json_body.get('stream', True) previous_messages = "\n\n".join([normalize_content(msg['content']) for msg in json_body['messages']]) input_tokens = calculate_tokens(previous_messages) msg_id = str(uuid.uuid4()) response_event = Event() response_text = [] total_output_tokens = 0 if not stream: return handle_non_stream(previous_messages, msg_id, model, input_tokens) log_request(request.remote_addr, request.path, 200) def generate(): nonlocal total_output_tokens start_event = create_event("message_start", { "type": "message_start", "message": { "id": msg_id, "type": "message", "role": "assistant", "model": model, "content": [], "stop_reason": None, "stop_sequence": None, "usage": {"input_tokens": input_tokens, "output_tokens": total_output_tokens}, }, }) logger.info("Sending message_start event", extra={ 'event_type': 'message_start', 'data': {'content': start_event} }) yield start_event block_start_event = create_event("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}) logger.info("Sending content_block_start event", extra={ 'event_type': 'content_block_start', 'data': {'content': block_start_event} }) yield block_start_event ping_event = create_event("ping", {"type": "ping"}) logger.info("Sending ping event", extra={ 'event_type': 'ping', 'data': {'content': ping_event} }) yield ping_event def on_query_progress(data): nonlocal total_output_tokens, response_text if 'text' in data: text = json.loads(data['text']) chunk = text['chunks'][-1] if text['chunks'] else None if chunk: response_text.append(chunk) chunk_tokens = calculate_tokens(chunk) total_output_tokens += chunk_tokens delta_event = create_event("content_block_delta", { "type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": chunk}, }) logger.info("Sending content_block_delta event", extra={ 'event_type': 'content_block_delta', 'data': { 'chunk': chunk, 'tokens': chunk_tokens, 'total_tokens': total_output_tokens, 'content': delta_event } }) yield delta_event if data.get('final', False): response_event.set() def on_connect(): logger.info("Connected to Perplexity AI", extra={'event_type': 'connection_established'}) emit_data = { "version": "2.9", "source": "default", "attachments": [], "language": "en-GB", "timezone": "Europe/London", "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.emit('perplexity_ask', (previous_messages, emit_data)) logger.info("Sent query to Perplexity AI", extra={ 'event_type': 'query_sent', 'data': { 'message': previous_messages[:100] + '...' if len(previous_messages) > 100 else previous_messages } }) sio.on('connect', on_connect) sio.on('query_progress', on_query_progress) try: sio.connect('wss://www.perplexity.ai/', **connect_opts, headers=sio_opts['extraHeaders']) while not response_event.is_set(): sio.sleep(0.1) except Exception as e: logger.error(f"Error during socket connection: {str(e)}", extra={ 'event_type': 'connection_error', 'data': {'error': 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)}"}, }) finally: if sio.connected: sio.disconnect() stop_event = create_event("content_block_stop", {"type": "content_block_stop", "index": 0}) logger.info("Sending content_block_stop event", extra={ 'event_type': 'content_block_stop', 'data': {'content': stop_event} }) yield stop_event message_delta_event = create_event("message_delta", { "type": "message_delta", "delta": {"stop_reason": "end_turn", "stop_sequence": None}, "usage": {"output_tokens": total_output_tokens}, }) logger.info("Sending message_delta event", extra={ 'event_type': 'message_delta', 'data': {'content': message_delta_event} }) yield message_delta_event message_stop_event = create_event("message_stop", {"type": "message_stop"}) logger.info("Sending message_stop event", extra={ 'event_type': 'message_stop', 'data': {'content': message_stop_event} }) yield message_stop_event return Response(generate(), content_type='text/event-stream') except Exception as e: logger.error(f"Request error: {str(e)}", extra={ 'event_type': 'request_error', 'data': {'error': str(e)} }) log_request(request.remote_addr, request.path, 400) return jsonify({"error": str(e)}), 400 def handle_non_stream(previous_messages, msg_id, model, input_tokens): try: response_event = Event() response_text = [] total_output_tokens = 0 def on_query_progress(data): nonlocal response_text, total_output_tokens if 'text' in data: text = json.loads(data['text']) chunk = text['chunks'][-1] if text['chunks'] else None if chunk: response_text.append(chunk) chunk_tokens = calculate_tokens(chunk) total_output_tokens += chunk_tokens if data.get('final', False): response_event.set() def on_connect(): logger.info("Connected to Perplexity AI (non-stream)", extra={'event_type': 'connection_established_non_stream'}) emit_data = { "version": "2.9", "source": "default", "attachments": [], "language": "en-GB", "timezone": "Europe/London", "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.emit('perplexity_ask', (previous_messages, emit_data)) sio.on('connect', on_connect) sio.on('query_progress', on_query_progress) sio.connect('wss://www.perplexity.ai/', **connect_opts, headers=sio_opts['extraHeaders']) response_event.wait(timeout=30) full_response = { "content": [{"text": ''.join(response_text), "type": "text"}], "id": msg_id, "model": model, "role": "assistant", "stop_reason": "end_turn", "stop_sequence": None, "type": "message", "usage": { "input_tokens": input_tokens, "output_tokens": total_output_tokens, }, } logger.info("Sending non-stream response", extra={ 'event_type': 'non_stream_response', 'data': {'content': full_response} }) return Response(json.dumps(full_response, ensure_ascii=False), content_type='application/json') except Exception as e: logger.error(f"Error during socket connection: {str(e)}", extra={ 'event_type': 'connection_error_non_stream', 'data': {'error': str(e)} }) return jsonify({"error": str(e)}), 500 finally: if sio.connected: sio.disconnect() @app.errorhandler(404) def not_found(error): log_request(request.remote_addr, request.path, 404) return "Not Found", 404 @app.errorhandler(500) def server_error(error): logger.error(f"Server error: {str(error)}", extra={ 'event_type': 'server_error', 'data': {'error': str(error)} }) log_request(request.remote_addr, request.path, 500) return "Something broke!", 500 if __name__ == '__main__': setup_logging() port = int(os.environ.get('PORT', 8081)) logger.info("Perplexity proxy starting", extra={ 'event_type': 'server_start', 'data': {'port': port} }) if not API_KEY: logger.warning("PPLX_KEY environment variable is not set", extra={'event_type': 'config_warning'}) app.run(host='0.0.0.0', port=port)