import os import json import uuid from datetime import datetime from flask import Flask, request, Response, jsonify, current_app import socketio import requests import logging from threading import Event app = Flask(__name__) logging.basicConfig(level=logging.INFO) # 从环境变量中获取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(ip, route, status): timestamp = datetime.now().isoformat() logging.info(f"{timestamp} - {ip} - {route} - {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 @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)", } } } }) @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(request.remote_addr, request.path, 200) 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} }) def generate(): 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', 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() try: sio.connect('wss://www.perplexity.ai/', **connect_opts, headers=sio_opts['extraHeaders']) 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)}"}, }) finally: if sio.connected: sio.disconnect() 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"}) log_request(request.remote_addr, request.path, 200) return Response(generate(), content_type='text/event-stream') except Exception as e: logging.error(f"Request error: {str(e)}") log_request(request.remote_addr, request.path, 400) return jsonify({"error": str(e)}), 400 @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): logging.error(f"Server error: {str(error)}") log_request(request.remote_addr, request.path, 500) return "Something broke!", 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)