Update app.py
Browse files
app.py
CHANGED
@@ -1,16 +1,43 @@
|
|
1 |
import os
|
2 |
import json
|
3 |
import uuid
|
4 |
-
import re
|
5 |
from datetime import datetime
|
6 |
from flask import Flask, request, Response, jsonify
|
7 |
import socketio
|
8 |
import requests
|
9 |
import logging
|
10 |
-
from threading import Event
|
|
|
11 |
|
12 |
app = Flask(__name__)
|
13 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
14 |
# 从环境变量中获取API密钥
|
15 |
API_KEY = os.environ.get('PPLX_KEY')
|
16 |
|
@@ -46,41 +73,15 @@ sio_opts = {
|
|
46 |
}
|
47 |
}
|
48 |
|
49 |
-
class CustomFormatter(logging.Formatter):
|
50 |
-
def format(self, record):
|
51 |
-
log_data = {
|
52 |
-
"timestamp": self.formatTime(record, self.datefmt),
|
53 |
-
"level": record.levelname,
|
54 |
-
"message": self.remove_ansi_escape(record.getMessage()),
|
55 |
-
}
|
56 |
-
if hasattr(record, 'event_type'):
|
57 |
-
log_data['event_type'] = record.event_type
|
58 |
-
if hasattr(record, 'data'):
|
59 |
-
log_data['data'] = record.data
|
60 |
-
return json.dumps(log_data, ensure_ascii=False, indent=2)
|
61 |
-
|
62 |
-
def remove_ansi_escape(self, text):
|
63 |
-
ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')
|
64 |
-
return ansi_escape.sub('', text)
|
65 |
-
|
66 |
-
def setup_logging():
|
67 |
-
logger = logging.getLogger()
|
68 |
-
logger.setLevel(logging.INFO)
|
69 |
-
handler = logging.StreamHandler()
|
70 |
-
handler.setFormatter(CustomFormatter())
|
71 |
-
logger.addHandler(handler)
|
72 |
-
|
73 |
-
logger = logging.getLogger(__name__)
|
74 |
-
|
75 |
def log_request(ip, route, status):
|
76 |
timestamp = datetime.now().isoformat()
|
77 |
-
logger.info(f"Request
|
78 |
-
'event_type': '
|
79 |
'data': {
|
80 |
-
'timestamp': timestamp,
|
81 |
'ip': ip,
|
82 |
'route': route,
|
83 |
-
'status': status
|
|
|
84 |
}
|
85 |
})
|
86 |
|
@@ -151,6 +152,7 @@ def messages():
|
|
151 |
|
152 |
msg_id = str(uuid.uuid4())
|
153 |
response_event = Event()
|
|
|
154 |
response_text = []
|
155 |
total_output_tokens = 0
|
156 |
|
@@ -241,10 +243,18 @@ def messages():
|
|
241 |
sio.on('connect', on_connect)
|
242 |
sio.on('query_progress', on_query_progress)
|
243 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
244 |
try:
|
245 |
sio.connect('wss://www.perplexity.ai/', **connect_opts, headers=sio_opts['extraHeaders'])
|
246 |
|
247 |
-
while not response_event.is_set():
|
248 |
sio.sleep(0.1)
|
249 |
while response_text:
|
250 |
chunk = response_text.pop(0)
|
@@ -254,6 +264,13 @@ def messages():
|
|
254 |
"delta": {"type": "text_delta", "text": chunk},
|
255 |
})
|
256 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
257 |
except Exception as e:
|
258 |
logger.error(f"Error during socket connection: {str(e)}", exc_info=True)
|
259 |
yield from send_event("content_block_delta", {
|
@@ -262,6 +279,7 @@ def messages():
|
|
262 |
"delta": {"type": "text_delta", "text": f"Error during socket connection: {str(e)}"},
|
263 |
})
|
264 |
finally:
|
|
|
265 |
if sio.connected:
|
266 |
sio.disconnect()
|
267 |
|
@@ -288,7 +306,7 @@ def handle_non_stream(previous_messages, msg_id, model, input_tokens):
|
|
288 |
total_output_tokens = 0
|
289 |
|
290 |
def on_query_progress(data):
|
291 |
-
nonlocal
|
292 |
if 'text' in data:
|
293 |
text = json.loads(data['text'])
|
294 |
chunk = text['chunks'][-1] if text['chunks'] else None
|
@@ -323,8 +341,13 @@ def handle_non_stream(previous_messages, msg_id, model, input_tokens):
|
|
323 |
|
324 |
sio.connect('wss://www.perplexity.ai/', **connect_opts, headers=sio_opts['extraHeaders'])
|
325 |
|
|
|
326 |
response_event.wait(timeout=30)
|
327 |
|
|
|
|
|
|
|
|
|
328 |
full_response = {
|
329 |
"content": [{"text": ''.join(response_text), "type": "text"}],
|
330 |
"id": msg_id,
|
@@ -345,10 +368,7 @@ def handle_non_stream(previous_messages, msg_id, model, input_tokens):
|
|
345 |
return Response(json.dumps(full_response, ensure_ascii=False), content_type='application/json')
|
346 |
|
347 |
except Exception as e:
|
348 |
-
logger.error(f"Error during socket connection: {str(e)}",
|
349 |
-
'event_type': 'connection_error_non_stream',
|
350 |
-
'data': {'error': str(e)}
|
351 |
-
})
|
352 |
return jsonify({"error": str(e)}), 500
|
353 |
finally:
|
354 |
if sio.connected:
|
@@ -361,17 +381,14 @@ def not_found(error):
|
|
361 |
|
362 |
@app.errorhandler(500)
|
363 |
def server_error(error):
|
364 |
-
logger.error(f"Server error: {str(error)}",
|
365 |
-
'event_type': 'server_error',
|
366 |
-
'data': {'error': str(error)}
|
367 |
-
})
|
368 |
log_request(request.remote_addr, request.path, 500)
|
369 |
return "Something broke!", 500
|
370 |
|
371 |
if __name__ == '__main__':
|
372 |
setup_logging()
|
373 |
port = int(os.environ.get('PORT', 8081))
|
374 |
-
logger.info("Perplexity proxy starting", extra={
|
375 |
'event_type': 'server_start',
|
376 |
'data': {'port': port}
|
377 |
})
|
|
|
1 |
import os
|
2 |
import json
|
3 |
import uuid
|
|
|
4 |
from datetime import datetime
|
5 |
from flask import Flask, request, Response, jsonify
|
6 |
import socketio
|
7 |
import requests
|
8 |
import logging
|
9 |
+
from threading import Event, Timer
|
10 |
+
import re
|
11 |
|
12 |
app = Flask(__name__)
|
13 |
|
14 |
+
# 自定义日志格式化器
|
15 |
+
class CustomFormatter(logging.Formatter):
|
16 |
+
def format(self, record):
|
17 |
+
log_data = {
|
18 |
+
"timestamp": self.formatTime(record, self.datefmt),
|
19 |
+
"level": record.levelname,
|
20 |
+
"message": self.remove_ansi_escape(record.getMessage()),
|
21 |
+
}
|
22 |
+
if hasattr(record, 'event_type'):
|
23 |
+
log_data['event_type'] = record.event_type
|
24 |
+
if hasattr(record, 'data'):
|
25 |
+
log_data['data'] = record.data
|
26 |
+
return json.dumps(log_data, ensure_ascii=False, indent=2)
|
27 |
+
|
28 |
+
def remove_ansi_escape(self, text):
|
29 |
+
ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')
|
30 |
+
return ansi_escape.sub('', text)
|
31 |
+
|
32 |
+
def setup_logging():
|
33 |
+
logger = logging.getLogger()
|
34 |
+
logger.setLevel(logging.INFO)
|
35 |
+
handler = logging.StreamHandler()
|
36 |
+
handler.setFormatter(CustomFormatter())
|
37 |
+
logger.addHandler(handler)
|
38 |
+
|
39 |
+
logger = logging.getLogger(__name__)
|
40 |
+
|
41 |
# 从环境变量中获取API密钥
|
42 |
API_KEY = os.environ.get('PPLX_KEY')
|
43 |
|
|
|
73 |
}
|
74 |
}
|
75 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
76 |
def log_request(ip, route, status):
|
77 |
timestamp = datetime.now().isoformat()
|
78 |
+
logger.info(f"Request logged", extra={
|
79 |
+
'event_type': 'request_log',
|
80 |
'data': {
|
|
|
81 |
'ip': ip,
|
82 |
'route': route,
|
83 |
+
'status': status,
|
84 |
+
'timestamp': timestamp
|
85 |
}
|
86 |
})
|
87 |
|
|
|
152 |
|
153 |
msg_id = str(uuid.uuid4())
|
154 |
response_event = Event()
|
155 |
+
timeout_event = Event()
|
156 |
response_text = []
|
157 |
total_output_tokens = 0
|
158 |
|
|
|
243 |
sio.on('connect', on_connect)
|
244 |
sio.on('query_progress', on_query_progress)
|
245 |
|
246 |
+
def timeout_handler():
|
247 |
+
logger.warning("Request timed out", extra={'event_type': 'request_timeout'})
|
248 |
+
timeout_event.set()
|
249 |
+
response_event.set()
|
250 |
+
|
251 |
+
timer = Timer(30, timeout_handler) # 30 seconds timeout
|
252 |
+
timer.start()
|
253 |
+
|
254 |
try:
|
255 |
sio.connect('wss://www.perplexity.ai/', **connect_opts, headers=sio_opts['extraHeaders'])
|
256 |
|
257 |
+
while not response_event.is_set() and not timeout_event.is_set():
|
258 |
sio.sleep(0.1)
|
259 |
while response_text:
|
260 |
chunk = response_text.pop(0)
|
|
|
264 |
"delta": {"type": "text_delta", "text": chunk},
|
265 |
})
|
266 |
|
267 |
+
if timeout_event.is_set():
|
268 |
+
yield from send_event("content_block_delta", {
|
269 |
+
"type": "content_block_delta",
|
270 |
+
"index": 0,
|
271 |
+
"delta": {"type": "text_delta", "text": "Request timed out"},
|
272 |
+
})
|
273 |
+
|
274 |
except Exception as e:
|
275 |
logger.error(f"Error during socket connection: {str(e)}", exc_info=True)
|
276 |
yield from send_event("content_block_delta", {
|
|
|
279 |
"delta": {"type": "text_delta", "text": f"Error during socket connection: {str(e)}"},
|
280 |
})
|
281 |
finally:
|
282 |
+
timer.cancel()
|
283 |
if sio.connected:
|
284 |
sio.disconnect()
|
285 |
|
|
|
306 |
total_output_tokens = 0
|
307 |
|
308 |
def on_query_progress(data):
|
309 |
+
nonlocal total_output_tokens, response_text
|
310 |
if 'text' in data:
|
311 |
text = json.loads(data['text'])
|
312 |
chunk = text['chunks'][-1] if text['chunks'] else None
|
|
|
341 |
|
342 |
sio.connect('wss://www.perplexity.ai/', **connect_opts, headers=sio_opts['extraHeaders'])
|
343 |
|
344 |
+
# Wait for response with timeout
|
345 |
response_event.wait(timeout=30)
|
346 |
|
347 |
+
if not response_text:
|
348 |
+
logger.warning("No response received (non-stream)", extra={'event_type': 'no_response_non_stream'})
|
349 |
+
return jsonify({"error": "No response received"}), 504
|
350 |
+
|
351 |
full_response = {
|
352 |
"content": [{"text": ''.join(response_text), "type": "text"}],
|
353 |
"id": msg_id,
|
|
|
368 |
return Response(json.dumps(full_response, ensure_ascii=False), content_type='application/json')
|
369 |
|
370 |
except Exception as e:
|
371 |
+
logger.error(f"Error during non-stream socket connection: {str(e)}", exc_info=True)
|
|
|
|
|
|
|
372 |
return jsonify({"error": str(e)}), 500
|
373 |
finally:
|
374 |
if sio.connected:
|
|
|
381 |
|
382 |
@app.errorhandler(500)
|
383 |
def server_error(error):
|
384 |
+
logger.error(f"Server error: {str(error)}", exc_info=True)
|
|
|
|
|
|
|
385 |
log_request(request.remote_addr, request.path, 500)
|
386 |
return "Something broke!", 500
|
387 |
|
388 |
if __name__ == '__main__':
|
389 |
setup_logging()
|
390 |
port = int(os.environ.get('PORT', 8081))
|
391 |
+
logger.info(f"Perplexity proxy starting", extra={
|
392 |
'event_type': 'server_start',
|
393 |
'data': {'port': port}
|
394 |
})
|