|
from flask import Flask, request, jsonify, Response |
|
from flask_socketio import SocketIO, emit |
|
from uuid import uuid4 |
|
import os |
|
import json |
|
import logging |
|
import requests |
|
from datetime import datetime |
|
|
|
app = Flask(__name__) |
|
socketio = SocketIO(app) |
|
|
|
|
|
API_KEY = os.getenv('PPLX_KEY') |
|
|
|
|
|
class ProxyAgent: |
|
def __init__(self): |
|
self.proxy = os.getenv('HTTP_PROXY') or os.getenv('HTTPS_PROXY') |
|
|
|
def get_session(self): |
|
session = requests.Session() |
|
if self.proxy: |
|
session.proxies = { |
|
'http': self.proxy, |
|
'https': self.proxy |
|
} |
|
return session |
|
|
|
agent = ProxyAgent() |
|
|
|
opts = { |
|
'agent': agent, |
|
'auth': { |
|
'jwt': 'anonymous-ask-user', |
|
}, |
|
'reconnection': False, |
|
'transports': ['websocket'], |
|
'path': '/socket.io', |
|
'hostname': 'www.perplexity.ai', |
|
'secure': True, |
|
'port': '443', |
|
'extraHeaders': { |
|
'Cookie': os.getenv('PPLX_COOKIE'), |
|
'User-Agent': os.getenv('USER_AGENT'), |
|
'Accept': '*/*', |
|
'priority': 'u=1, i', |
|
'Referer': 'https://www.perplexity.ai/', |
|
}, |
|
} |
|
|
|
|
|
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 log_request(req, status): |
|
timestamp = datetime.now().isoformat() |
|
ip = req.remote_addr |
|
route = req.path |
|
logging.info(f'{timestamp} - {ip} - {route} - {status}') |
|
|
|
|
|
@app.route('/', methods=['GET']) |
|
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)', |
|
|
|
} |
|
} |
|
} |
|
}) |
|
|
|
|
|
@app.route('/ai/v1/messages', methods=['POST']) |
|
def ai_messages(): |
|
validation_error = validate_api_key() |
|
if validation_error: |
|
return validation_error |
|
|
|
raw_body = request.data.decode('utf-8') |
|
try: |
|
json_body = json.loads(raw_body) |
|
if json_body.get('stream') is False: |
|
log_request(request, 200) |
|
return jsonify({ |
|
'id': str(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, |
|
}, |
|
}) |
|
elif json_body.get('stream') is True: |
|
def generate(): |
|
try: |
|
|
|
user_message = [{'question': '', 'answer': ''}] |
|
last_update = True |
|
if json_body.get('system'): |
|
json_body['messages'].insert(0, {'role': 'system', 'content': json_body['system']}) |
|
for msg in json_body['messages']: |
|
if msg['role'] in ['system', 'user']: |
|
if last_update: |
|
user_message[-1]['question'] += msg['content'] + '\n' |
|
elif user_message[-1]['question'] == '': |
|
user_message[-1]['question'] += msg['content'] + '\n' |
|
else: |
|
user_message.append({'question': msg['content'] + '\n', 'answer': ''}) |
|
last_update = True |
|
elif msg['role'] == 'assistant': |
|
if not last_update: |
|
user_message[-1]['answer'] += msg['content'] + '\n' |
|
elif user_message[-1]['answer'] == '': |
|
user_message[-1]['answer'] += msg['content'] + '\n' |
|
else: |
|
user_message.append({'question': '', 'answer': msg['content'] + '\n'}) |
|
last_update = False |
|
|
|
previous_messages = '\n\n'.join(msg['content'] for msg in json_body['messages']) |
|
msg_id = str(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'}) |
|
|
|
|
|
socket = socketio.connect('wss://www.perplexity.ai/', **opts) |
|
|
|
@socket.on('connect') |
|
def on_connect(): |
|
logging.info(' > [Connected]') |
|
socket.emit('perplexity_ask', previous_messages, { |
|
'version': '2.9', |
|
'source': 'default', |
|
'attachments': [], |
|
'language': 'en-GB', |
|
'timezone': 'Europe/London', |
|
'search_focus': 'writing', |
|
'frontend_uuid': str(uuid4()), |
|
'mode': 'concise', |
|
'is_related_query': False, |
|
'is_default_related_query': False, |
|
'visitor_id': str(uuid4()), |
|
'frontend_context_uuid': str(uuid4()), |
|
'prompt_source': 'user', |
|
'query_source': 'home' |
|
}, callback=on_response) |
|
|
|
def on_response(response): |
|
logging.info(response) |
|
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) |
|
|
|
@socket.on('query_progress') |
|
def on_query_progress(data): |
|
if data.get('text'): |
|
text = json.loads(data['text']) |
|
chunk = text['chunks'][-1] |
|
if chunk: |
|
yield create_event('content_block_delta', { |
|
'type': 'content_block_delta', |
|
'index': 0, |
|
'delta': {'type': 'text_delta', 'text': chunk}, |
|
}) |
|
|
|
@socket.on('disconnect') |
|
def on_disconnect(): |
|
logging.info(' > [Disconnected]') |
|
|
|
@socket.on('error') |
|
def on_error(error): |
|
yield create_event('content_block_delta', { |
|
'type': 'content_block_delta', |
|
'index': 0, |
|
'delta': {'type': 'text_delta', 'text': 'Error occured 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'}) |
|
logging.error(error) |
|
log_request(request, 500) |
|
|
|
@socket.on('connect_error') |
|
def on_connect_error(error): |
|
yield create_event('content_block_delta', { |
|
'type': 'content_block_delta', |
|
'index': 0, |
|
'delta': {'type': 'text_delta', 'text': 'Failed to connect to the Perplexity.ai 连接到Perplexity失败\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'}) |
|
logging.error(error) |
|
log_request(request, 500) |
|
|
|
@socket.on('close') |
|
def on_close(): |
|
logging.info(' > [Client closed]') |
|
socket.disconnect() |
|
|
|
except Exception as e: |
|
logging.error(e) |
|
yield create_event('content_block_delta', { |
|
'type': 'content_block_delta', |
|
'index': 0, |
|
'delta': {'type': 'text_delta', 'text': 'Error occured 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) |
|
|
|
return Response(generate(), mimetype='text/event-stream') |
|
else: |
|
raise ValueError('Invalid request') |
|
except Exception as e: |
|
logging.error(e) |
|
return jsonify({'error': str(e)}), 400 |
|
|
|
|
|
@app.errorhandler(404) |
|
def not_found(e): |
|
log_request(request, 404) |
|
return 'Not Found', 404 |
|
|
|
|
|
@app.errorhandler(500) |
|
def internal_server_error(e): |
|
logging.error(e) |
|
log_request(request, 500) |
|
return 'Something broke!', 500 |
|
|
|
if __name__ == '__main__': |
|
port = int(os.getenv('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.') |
|
socketio.run(app, port=port) |
|
|
|
|
|
def create_event(event, data): |
|
if isinstance(data, dict): |
|
data = json.dumps(data) |
|
return f'event: {event}\ndata: {data}\n\n' |