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 app = Flask(__name__) # 从环境变量中获取日志级别,默认为INFO log_level = os.environ.get('LOG_LEVEL', 'INFO').upper() logging.basicConfig(level=getattr(logging, log_level)) # 从环境变量中获取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'], # 允许回退到轮询 } # 其他选项 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/', } } def log_request(timestamp, ip, route, status, request_id): logging.info(f"{timestamp} - {ip} - {route} - {status} - Request ID: {request_id}") 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 return None @app.route('/') def root(): log_request(datetime.now().isoformat(), request.remote_addr, request.path, 200, str(uuid.uuid4())) 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)", } } } }) @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 if not json_body.get('stream', False): log_request(datetime.now().isoformat(), request.remote_addr, request.path, 200, str(uuid.uuid4())) return jsonify({ "id": str(uuid.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}, }, }) 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}, }) @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"}) return Response(generate(), content_type='text/event-stream') 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 @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 @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 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)