File size: 6,448 Bytes
e3596d7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 |
import os
import json
import uuid
from flask import Flask, request, Response, jsonify, stream_with_context
from flask_limiter import Limiter
from flask_limiter.util import get_remote_address
import socketio
import requests
from requests.exceptions import RequestException
app = Flask(__name__)
# 从环境变量中获取配置
API_KEY = os.environ.get('PPLX_KEY')
PPLX_COOKIE = os.environ.get('PPLX_COOKIE')
USER_AGENT = os.environ.get('USER_AGENT')
PROXY = os.environ.get('PROXY')
# 设置限流
limiter = Limiter(
get_remote_address,
app=app,
default_limits=["100 per 5 minutes"],
storage_uri="memory://"
)
# Socket.IO 客户端设置
sio = socketio.Client()
# 代理设置
proxies = {'https': PROXY} if PROXY else None
@app.route('/')
def home():
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('/health')
def health_check():
return jsonify({"status": "OK"}), 200
def validate_api_key():
api_key = request.headers.get('x-api-key')
if api_key != API_KEY:
return jsonify({"error": "Invalid API key"}), 401
@app.route('/ai/v1/messages', methods=['POST'])
@limiter.limit("100 per 15 minutes")
def ai_messages():
auth_result = validate_api_key()
if auth_result:
return auth_result
try:
data = request.json
if not data.get('stream', False):
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 data['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"})
try:
sio.connect('wss://www.perplexity.ai', transports=['websocket'],
headers={
'Cookie': PPLX_COOKIE,
'User-Agent': USER_AGENT
},
http_session=requests.Session() if PROXY else None,
proxies=proxies)
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.on('query_progress')
def on_query_progress(data):
if data.get('text'):
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},
})
sio.wait()
except Exception as e:
print(f"Socket error: {e}")
yield create_event("content_block_delta", {
"type": "content_block_delta",
"index": 0,
"delta": {"type": "text_delta", "text": "An error occurred while processing your request"},
})
finally:
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"})
return Response(stream_with_context(generate()), content_type='text/event-stream')
except Exception as e:
print(f"Request error: {e}")
return jsonify({"error": str(e)}), 400
def create_event(event, data):
if isinstance(data, dict):
data = json.dumps(data)
return f"event: {event}\ndata: {data}\n\n"
@app.errorhandler(404)
def not_found(error):
return jsonify({"error": "Not Found"}), 404
@app.errorhandler(500)
def server_error(error):
return jsonify({"error": "Internal Server Error"}), 500
if __name__ == '__main__':
port = int(os.environ.get('PORT', 8081))
app.run(host='0.0.0.0', port=port)
|