import os | |
import json | |
import uuid | |
from flask import Flask, request, Response, jsonify | |
import socketio | |
from aiohttp import web | |
from aiohttp_socks import ProxyConnector | |
app = Flask(__name__) | |
sio = socketio.Client() | |
# 从环境变量中获取API密钥 | |
API_KEY = os.environ.get('PPLX_KEY') | |
# 代理设置 | |
proxy_url = os.environ.get('PROXY_URL') | |
connector = ProxyConnector.from_url(proxy_url) if proxy_url else None | |
sio_opts = { | |
'auth': { | |
'jwt': 'anonymous-ask-user', | |
}, | |
'reconnection': False, | |
'transports': ['websocket'], | |
'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(req, status): | |
timestamp = datetime.now().isoformat() | |
ip = req.remote_addr | |
route = req.path | |
print(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, 401) | |
return jsonify({"error": "Invalid API key"}), 401 | |
return None | |
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)", | |
} | |
} | |
} | |
}) | |
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, 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()) | |
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"}) | |
def on_connect(): | |
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" | |
}) | |
def on_query_progress(data): | |
if 'text' in data: | |
text = json.loads(data['text']) | |
chunk = text['chunks'][-1] if text['chunks'] else None | |
if chunk: | |
yield create_event("content_block_delta", { | |
"type": "content_block_delta", | |
"index": 0, | |
"delta": {"type": "text_delta", "text": chunk}, | |
}) | |
def on_disconnect(): | |
print(" > [Disconnected]") | |
def on_error(error): | |
print(f"Socket error: {error}") | |
yield create_event("content_block_delta", { | |
"type": "content_block_delta", | |
"index": 0, | |
"delta": {"type": "text_delta", "text": "Error occurred 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) | |
sio.connect('wss://www.perplexity.ai/', **sio_opts) | |
sio.wait() | |
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": 12}, | |
}) | |
yield create_event("message_stop", {"type": "message_stop"}) | |
log_request(request, 200) | |
return Response(generate(), content_type='text/event-stream') | |
except Exception as e: | |
print(f"Request error: {str(e)}") | |
log_request(request, 400) | |
return jsonify({"error": str(e)}), 400 | |
def not_found(error): | |
log_request(request, 404) | |
return "Not Found", 404 | |
def server_error(error): | |
print(f"Server error: {str(error)}") | |
log_request(request, 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)) | |
print(f"Perplexity proxy listening on port {port}") | |
if not API_KEY: | |
print("Warning: PPLX_KEY environment variable is not set. API key validation will fail.") | |
app.run(host='0.0.0.0', port=port) | |