Update app.py
Browse files
app.py
CHANGED
@@ -8,11 +8,9 @@ import requests
|
|
8 |
import logging
|
9 |
from threading import Event, Timer
|
10 |
import re
|
11 |
-
import asyncio
|
12 |
|
13 |
app = Flask(__name__)
|
14 |
|
15 |
-
# 自定义日志格式化器
|
16 |
class CustomFormatter(logging.Formatter):
|
17 |
def format(self, record):
|
18 |
log_data = {
|
@@ -39,13 +37,9 @@ def setup_logging():
|
|
39 |
|
40 |
logger = logging.getLogger(__name__)
|
41 |
|
42 |
-
# 从环境变量中获取API密钥
|
43 |
API_KEY = os.environ.get('PPLX_KEY')
|
44 |
-
|
45 |
-
# 代理设置
|
46 |
proxy_url = os.environ.get('PROXY_URL')
|
47 |
|
48 |
-
# 设置代理
|
49 |
if proxy_url:
|
50 |
proxies = {
|
51 |
'http': proxy_url,
|
@@ -58,12 +52,10 @@ else:
|
|
58 |
|
59 |
sio = socketio.Client(http_session=transport, logger=False, engineio_logger=False)
|
60 |
|
61 |
-
# 连接选项
|
62 |
connect_opts = {
|
63 |
'transports': ['websocket', 'polling'],
|
64 |
}
|
65 |
|
66 |
-
# 其他选项
|
67 |
sio_opts = {
|
68 |
'extraHeaders': {
|
69 |
'Cookie': os.environ.get('PPLX_COOKIE'),
|
@@ -115,6 +107,11 @@ def create_event(event, data):
|
|
115 |
data = json.dumps(data, ensure_ascii=False)
|
116 |
return f"event: {event}\ndata: {data}\n\n"
|
117 |
|
|
|
|
|
|
|
|
|
|
|
118 |
@app.route('/')
|
119 |
def root():
|
120 |
log_request(request.remote_addr, request.path, 200)
|
@@ -130,18 +127,13 @@ def root():
|
|
130 |
},
|
131 |
"body": {
|
132 |
"messages": "Array of message objects",
|
133 |
-
"stream": "Boolean (
|
134 |
-
"model": "Model to be used (optional, defaults to claude-3-
|
135 |
}
|
136 |
}
|
137 |
}
|
138 |
})
|
139 |
|
140 |
-
async def process_large_request(previous_messages, model, input_tokens):
|
141 |
-
# 这个函数用于异步处理大型请求
|
142 |
-
# 实际实现时,你可能需要将这个过程放到后台任务队列中
|
143 |
-
pass
|
144 |
-
|
145 |
@app.route('/ai/v1/messages', methods=['POST'])
|
146 |
def messages():
|
147 |
auth_error = validate_api_key()
|
@@ -150,14 +142,12 @@ def messages():
|
|
150 |
|
151 |
try:
|
152 |
json_body = request.json
|
153 |
-
model = json_body.get('model', 'claude-3-
|
154 |
-
stream = json_body.get('stream',
|
155 |
|
156 |
previous_messages = "\n\n".join([normalize_content(msg['content']) for msg in json_body['messages']])
|
157 |
input_tokens = calculate_tokens(previous_messages)
|
158 |
-
|
159 |
-
# 根据 input_tokens 动态调整超时时间
|
160 |
-
timeout = max(30, min(300, input_tokens // 1000 * 30)) # 最少30秒,最多300秒
|
161 |
|
162 |
msg_id = str(uuid.uuid4())
|
163 |
response_event = Event()
|
@@ -165,14 +155,6 @@ def messages():
|
|
165 |
response_text = []
|
166 |
total_output_tokens = 0
|
167 |
|
168 |
-
if input_tokens > 100000: # 如果 tokens 数量特别大,使用异步处理
|
169 |
-
task_id = str(uuid.uuid4())
|
170 |
-
asyncio.create_task(process_large_request(previous_messages, model, input_tokens))
|
171 |
-
return jsonify({
|
172 |
-
"message": "Request is being processed asynchronously",
|
173 |
-
"task_id": task_id
|
174 |
-
}), 202
|
175 |
-
|
176 |
if not stream:
|
177 |
return handle_non_stream(previous_messages, msg_id, model, input_tokens, timeout)
|
178 |
|
@@ -189,7 +171,6 @@ def messages():
|
|
189 |
})
|
190 |
yield event
|
191 |
|
192 |
-
# Send initial events
|
193 |
yield from send_event("message_start", {
|
194 |
"type": "message_start",
|
195 |
"message": {
|
@@ -223,12 +204,6 @@ def messages():
|
|
223 |
'total_tokens': total_output_tokens
|
224 |
}
|
225 |
})
|
226 |
-
# 发送进度更新
|
227 |
-
yield from send_event("progress", {
|
228 |
-
"type": "progress",
|
229 |
-
"processed_tokens": total_output_tokens,
|
230 |
-
"total_tokens": input_tokens
|
231 |
-
})
|
232 |
|
233 |
if data.get('final', False):
|
234 |
logger.info("Final response received", extra={
|
@@ -267,11 +242,17 @@ def messages():
|
|
267 |
sio.on('query_progress', on_query_progress)
|
268 |
|
269 |
def timeout_handler():
|
270 |
-
logger.warning("Request timed out", extra={
|
|
|
|
|
|
|
|
|
|
|
|
|
271 |
timeout_event.set()
|
272 |
response_event.set()
|
273 |
|
274 |
-
timer = Timer(timeout, timeout_handler)
|
275 |
timer.start()
|
276 |
|
277 |
try:
|
@@ -291,7 +272,7 @@ def messages():
|
|
291 |
yield from send_event("content_block_delta", {
|
292 |
"type": "content_block_delta",
|
293 |
"index": 0,
|
294 |
-
"delta": {"type": "text_delta", "text": "Request timed out
|
295 |
})
|
296 |
|
297 |
except Exception as e:
|
@@ -306,7 +287,6 @@ def messages():
|
|
306 |
if sio.connected:
|
307 |
sio.disconnect()
|
308 |
|
309 |
-
# Send final events
|
310 |
yield from send_event("content_block_stop", {"type": "content_block_stop", "index": 0})
|
311 |
yield from send_event("message_delta", {
|
312 |
"type": "message_delta",
|
@@ -364,11 +344,16 @@ def handle_non_stream(previous_messages, msg_id, model, input_tokens, timeout):
|
|
364 |
|
365 |
sio.connect('wss://www.perplexity.ai/', **connect_opts, headers=sio_opts['extraHeaders'])
|
366 |
|
367 |
-
# Wait for response with timeout
|
368 |
response_event.wait(timeout=timeout)
|
369 |
|
370 |
if not response_text:
|
371 |
-
logger.warning("No response received (non-stream)", extra={
|
|
|
|
|
|
|
|
|
|
|
|
|
372 |
return jsonify({"error": "No response received"}), 504
|
373 |
|
374 |
full_response = {
|
|
|
8 |
import logging
|
9 |
from threading import Event, Timer
|
10 |
import re
|
|
|
11 |
|
12 |
app = Flask(__name__)
|
13 |
|
|
|
14 |
class CustomFormatter(logging.Formatter):
|
15 |
def format(self, record):
|
16 |
log_data = {
|
|
|
37 |
|
38 |
logger = logging.getLogger(__name__)
|
39 |
|
|
|
40 |
API_KEY = os.environ.get('PPLX_KEY')
|
|
|
|
|
41 |
proxy_url = os.environ.get('PROXY_URL')
|
42 |
|
|
|
43 |
if proxy_url:
|
44 |
proxies = {
|
45 |
'http': proxy_url,
|
|
|
52 |
|
53 |
sio = socketio.Client(http_session=transport, logger=False, engineio_logger=False)
|
54 |
|
|
|
55 |
connect_opts = {
|
56 |
'transports': ['websocket', 'polling'],
|
57 |
}
|
58 |
|
|
|
59 |
sio_opts = {
|
60 |
'extraHeaders': {
|
61 |
'Cookie': os.environ.get('PPLX_COOKIE'),
|
|
|
107 |
data = json.dumps(data, ensure_ascii=False)
|
108 |
return f"event: {event}\ndata: {data}\n\n"
|
109 |
|
110 |
+
def calculate_timeout(input_tokens):
|
111 |
+
base_timeout = 30
|
112 |
+
additional_time = min(max(input_tokens // 100 * 5, 0), 270)
|
113 |
+
return base_timeout + additional_time
|
114 |
+
|
115 |
@app.route('/')
|
116 |
def root():
|
117 |
log_request(request.remote_addr, request.path, 200)
|
|
|
127 |
},
|
128 |
"body": {
|
129 |
"messages": "Array of message objects",
|
130 |
+
"stream": "Boolean (optional, defaults to false)",
|
131 |
+
"model": "Model to be used (optional, defaults to claude-3-5-sonnet-20240620)"
|
132 |
}
|
133 |
}
|
134 |
}
|
135 |
})
|
136 |
|
|
|
|
|
|
|
|
|
|
|
137 |
@app.route('/ai/v1/messages', methods=['POST'])
|
138 |
def messages():
|
139 |
auth_error = validate_api_key()
|
|
|
142 |
|
143 |
try:
|
144 |
json_body = request.json
|
145 |
+
model = json_body.get('model', 'claude-3-5-sonnet-20240620')
|
146 |
+
stream = json_body.get('stream', False)
|
147 |
|
148 |
previous_messages = "\n\n".join([normalize_content(msg['content']) for msg in json_body['messages']])
|
149 |
input_tokens = calculate_tokens(previous_messages)
|
150 |
+
timeout = calculate_timeout(input_tokens)
|
|
|
|
|
151 |
|
152 |
msg_id = str(uuid.uuid4())
|
153 |
response_event = Event()
|
|
|
155 |
response_text = []
|
156 |
total_output_tokens = 0
|
157 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
158 |
if not stream:
|
159 |
return handle_non_stream(previous_messages, msg_id, model, input_tokens, timeout)
|
160 |
|
|
|
171 |
})
|
172 |
yield event
|
173 |
|
|
|
174 |
yield from send_event("message_start", {
|
175 |
"type": "message_start",
|
176 |
"message": {
|
|
|
204 |
'total_tokens': total_output_tokens
|
205 |
}
|
206 |
})
|
|
|
|
|
|
|
|
|
|
|
|
|
207 |
|
208 |
if data.get('final', False):
|
209 |
logger.info("Final response received", extra={
|
|
|
242 |
sio.on('query_progress', on_query_progress)
|
243 |
|
244 |
def timeout_handler():
|
245 |
+
logger.warning("Request timed out", extra={
|
246 |
+
'event_type': 'request_timeout',
|
247 |
+
'data': {
|
248 |
+
'input_tokens': input_tokens,
|
249 |
+
'timeout': timeout
|
250 |
+
}
|
251 |
+
})
|
252 |
timeout_event.set()
|
253 |
response_event.set()
|
254 |
|
255 |
+
timer = Timer(timeout, timeout_handler)
|
256 |
timer.start()
|
257 |
|
258 |
try:
|
|
|
272 |
yield from send_event("content_block_delta", {
|
273 |
"type": "content_block_delta",
|
274 |
"index": 0,
|
275 |
+
"delta": {"type": "text_delta", "text": "Request timed out"},
|
276 |
})
|
277 |
|
278 |
except Exception as e:
|
|
|
287 |
if sio.connected:
|
288 |
sio.disconnect()
|
289 |
|
|
|
290 |
yield from send_event("content_block_stop", {"type": "content_block_stop", "index": 0})
|
291 |
yield from send_event("message_delta", {
|
292 |
"type": "message_delta",
|
|
|
344 |
|
345 |
sio.connect('wss://www.perplexity.ai/', **connect_opts, headers=sio_opts['extraHeaders'])
|
346 |
|
|
|
347 |
response_event.wait(timeout=timeout)
|
348 |
|
349 |
if not response_text:
|
350 |
+
logger.warning("No response received (non-stream)", extra={
|
351 |
+
'event_type': 'no_response_non_stream',
|
352 |
+
'data': {
|
353 |
+
'input_tokens': input_tokens,
|
354 |
+
'timeout': timeout
|
355 |
+
}
|
356 |
+
})
|
357 |
return jsonify({"error": "No response received"}), 504
|
358 |
|
359 |
full_response = {
|