Update app.py
Browse files
app.py
CHANGED
@@ -2,13 +2,15 @@ import os
|
|
2 |
import json
|
3 |
import uuid
|
4 |
import time
|
5 |
-
|
6 |
-
|
7 |
import socketio
|
8 |
import requests
|
9 |
import logging
|
10 |
-
from threading import Event
|
11 |
import re
|
|
|
|
|
|
|
12 |
|
13 |
app = Flask(__name__)
|
14 |
|
@@ -48,8 +50,7 @@ if proxy_url:
|
|
48 |
else:
|
49 |
transport = None
|
50 |
|
51 |
-
sio = socketio.
|
52 |
-
|
53 |
connect_opts = {'transports': ['websocket', 'polling']}
|
54 |
|
55 |
sio_opts = {
|
@@ -125,7 +126,7 @@ def root():
|
|
125 |
})
|
126 |
|
127 |
@app.route('/ai/v1/messages', methods=['POST'])
|
128 |
-
def messages():
|
129 |
auth_error = validate_api_key()
|
130 |
if auth_error:
|
131 |
return auth_error
|
@@ -139,20 +140,20 @@ def messages():
|
|
139 |
input_tokens = calculate_tokens(previous_messages)
|
140 |
|
141 |
msg_id = str(uuid.uuid4())
|
142 |
-
response_event = Event()
|
143 |
response_text = []
|
144 |
total_output_tokens = 0
|
145 |
|
146 |
if not stream:
|
147 |
-
return handle_non_stream(previous_messages, msg_id, model, input_tokens)
|
148 |
|
149 |
log_request(request.remote_addr, request.path, 200)
|
150 |
|
151 |
-
def generate():
|
152 |
nonlocal total_output_tokens
|
153 |
start_time = time.time()
|
154 |
last_activity_time = start_time
|
155 |
-
timeout = max(
|
156 |
|
157 |
def send_event(event_type, data):
|
158 |
event = create_event(event_type, data)
|
@@ -160,9 +161,9 @@ def messages():
|
|
160 |
'event_type': event_type,
|
161 |
'data': {'content': event}
|
162 |
})
|
163 |
-
|
164 |
|
165 |
-
yield
|
166 |
"type": "message_start",
|
167 |
"message": {
|
168 |
"id": msg_id,
|
@@ -175,10 +176,11 @@ def messages():
|
|
175 |
"usage": {"input_tokens": input_tokens, "output_tokens": total_output_tokens},
|
176 |
},
|
177 |
})
|
178 |
-
yield
|
179 |
-
yield
|
180 |
|
181 |
-
|
|
|
182 |
nonlocal total_output_tokens, response_text, last_activity_time
|
183 |
last_activity_time = time.time()
|
184 |
if 'text' in data:
|
@@ -206,7 +208,8 @@ def messages():
|
|
206 |
})
|
207 |
response_event.set()
|
208 |
|
209 |
-
|
|
|
210 |
logger.info("Connected to Perplexity AI", extra={'event_type': 'connection_established'})
|
211 |
emit_data = {
|
212 |
"version": "2.9",
|
@@ -222,7 +225,7 @@ def messages():
|
|
222 |
"prompt_source": "user",
|
223 |
"query_source": "home"
|
224 |
}
|
225 |
-
sio.emit('perplexity_ask', (previous_messages, emit_data))
|
226 |
logger.info("Sent query to Perplexity AI", extra={
|
227 |
'event_type': 'query_sent',
|
228 |
'data': {
|
@@ -230,20 +233,23 @@ def messages():
|
|
230 |
}
|
231 |
})
|
232 |
|
233 |
-
|
234 |
-
|
|
|
|
|
235 |
|
236 |
try:
|
237 |
-
sio.connect('wss://www.perplexity.ai/', **connect_opts, headers=sio_opts['extraHeaders'])
|
|
|
238 |
|
239 |
while not response_event.is_set() and (time.time() - start_time) < timeout:
|
240 |
current_time = time.time()
|
241 |
-
if current_time - last_activity_time >
|
242 |
-
logger.warning("No activity for
|
243 |
-
|
244 |
while response_text:
|
245 |
chunk = response_text.pop(0)
|
246 |
-
yield
|
247 |
"type": "content_block_delta",
|
248 |
"index": 0,
|
249 |
"delta": {"type": "text_delta", "text": chunk},
|
@@ -259,7 +265,7 @@ def messages():
|
|
259 |
'elapsed_time': time.time() - start_time
|
260 |
}
|
261 |
})
|
262 |
-
yield
|
263 |
"type": "content_block_delta",
|
264 |
"index": 0,
|
265 |
"delta": {"type": "text_delta", "text": f"Request timed out after {timeout} seconds"},
|
@@ -267,22 +273,23 @@ def messages():
|
|
267 |
|
268 |
except Exception as e:
|
269 |
logger.error(f"Error during socket connection: {str(e)}", exc_info=True)
|
270 |
-
yield
|
271 |
"type": "content_block_delta",
|
272 |
"index": 0,
|
273 |
"delta": {"type": "text_delta", "text": f"Error during socket connection: {str(e)}"},
|
274 |
})
|
275 |
finally:
|
|
|
276 |
if sio.connected:
|
277 |
-
sio.disconnect()
|
278 |
|
279 |
-
yield
|
280 |
-
yield
|
281 |
"type": "message_delta",
|
282 |
"delta": {"stop_reason": "end_turn", "stop_sequence": None},
|
283 |
"usage": {"output_tokens": total_output_tokens},
|
284 |
})
|
285 |
-
yield
|
286 |
|
287 |
return Response(generate(), content_type='text/event-stream')
|
288 |
|
@@ -406,6 +413,6 @@ if __name__ == '__main__':
|
|
406 |
})
|
407 |
if not API_KEY:
|
408 |
logger.warning("PPLX_KEY environment variable is not set", extra={'event_type': 'config_warning'})
|
409 |
-
app.run(host='0.0.0.0', port=port)
|
410 |
|
411 |
|
|
|
2 |
import json
|
3 |
import uuid
|
4 |
import time
|
5 |
+
import asyncio
|
6 |
+
import socketio
|
7 |
import socketio
|
8 |
import requests
|
9 |
import logging
|
|
|
10 |
import re
|
11 |
+
from datetime import datetime
|
12 |
+
from flask import Flask, request, Response, jsonify
|
13 |
+
from threading import Event
|
14 |
|
15 |
app = Flask(__name__)
|
16 |
|
|
|
50 |
else:
|
51 |
transport = None
|
52 |
|
53 |
+
sio = socketio.AsyncClient(http_session=transport, logger=False, engineio_logger=False)
|
|
|
54 |
connect_opts = {'transports': ['websocket', 'polling']}
|
55 |
|
56 |
sio_opts = {
|
|
|
126 |
})
|
127 |
|
128 |
@app.route('/ai/v1/messages', methods=['POST'])
|
129 |
+
async def messages():
|
130 |
auth_error = validate_api_key()
|
131 |
if auth_error:
|
132 |
return auth_error
|
|
|
140 |
input_tokens = calculate_tokens(previous_messages)
|
141 |
|
142 |
msg_id = str(uuid.uuid4())
|
143 |
+
response_event = asyncio.Event()
|
144 |
response_text = []
|
145 |
total_output_tokens = 0
|
146 |
|
147 |
if not stream:
|
148 |
+
return await handle_non_stream(previous_messages, msg_id, model, input_tokens)
|
149 |
|
150 |
log_request(request.remote_addr, request.path, 200)
|
151 |
|
152 |
+
async def generate():
|
153 |
nonlocal total_output_tokens
|
154 |
start_time = time.time()
|
155 |
last_activity_time = start_time
|
156 |
+
timeout = max(300, input_tokens / 100) # 动态设置超时时间,最少300秒
|
157 |
|
158 |
def send_event(event_type, data):
|
159 |
event = create_event(event_type, data)
|
|
|
161 |
'event_type': event_type,
|
162 |
'data': {'content': event}
|
163 |
})
|
164 |
+
return event
|
165 |
|
166 |
+
yield send_event("message_start", {
|
167 |
"type": "message_start",
|
168 |
"message": {
|
169 |
"id": msg_id,
|
|
|
176 |
"usage": {"input_tokens": input_tokens, "output_tokens": total_output_tokens},
|
177 |
},
|
178 |
})
|
179 |
+
yield send_event("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}})
|
180 |
+
yield send_event("ping", {"type": "ping"})
|
181 |
|
182 |
+
@sio.event
|
183 |
+
async def query_progress(data):
|
184 |
nonlocal total_output_tokens, response_text, last_activity_time
|
185 |
last_activity_time = time.time()
|
186 |
if 'text' in data:
|
|
|
208 |
})
|
209 |
response_event.set()
|
210 |
|
211 |
+
@sio.event
|
212 |
+
async def connect():
|
213 |
logger.info("Connected to Perplexity AI", extra={'event_type': 'connection_established'})
|
214 |
emit_data = {
|
215 |
"version": "2.9",
|
|
|
225 |
"prompt_source": "user",
|
226 |
"query_source": "home"
|
227 |
}
|
228 |
+
await sio.emit('perplexity_ask', (previous_messages, emit_data))
|
229 |
logger.info("Sent query to Perplexity AI", extra={
|
230 |
'event_type': 'query_sent',
|
231 |
'data': {
|
|
|
233 |
}
|
234 |
})
|
235 |
|
236 |
+
async def heartbeat():
|
237 |
+
while not response_event.is_set():
|
238 |
+
await sio.emit('ping')
|
239 |
+
await asyncio.sleep(25)
|
240 |
|
241 |
try:
|
242 |
+
await sio.connect('wss://www.perplexity.ai/', **connect_opts, headers=sio_opts['extraHeaders'])
|
243 |
+
heartbeat_task = asyncio.create_task(heartbeat())
|
244 |
|
245 |
while not response_event.is_set() and (time.time() - start_time) < timeout:
|
246 |
current_time = time.time()
|
247 |
+
if current_time - last_activity_time > 60: # 如果60秒内没有活动,记录警告
|
248 |
+
logger.warning("No activity for 60 seconds", extra={'event_type': 'inactivity_warning'})
|
249 |
+
await asyncio.sleep(1)
|
250 |
while response_text:
|
251 |
chunk = response_text.pop(0)
|
252 |
+
yield send_event("content_block_delta", {
|
253 |
"type": "content_block_delta",
|
254 |
"index": 0,
|
255 |
"delta": {"type": "text_delta", "text": chunk},
|
|
|
265 |
'elapsed_time': time.time() - start_time
|
266 |
}
|
267 |
})
|
268 |
+
yield send_event("content_block_delta", {
|
269 |
"type": "content_block_delta",
|
270 |
"index": 0,
|
271 |
"delta": {"type": "text_delta", "text": f"Request timed out after {timeout} seconds"},
|
|
|
273 |
|
274 |
except Exception as e:
|
275 |
logger.error(f"Error during socket connection: {str(e)}", exc_info=True)
|
276 |
+
yield send_event("content_block_delta", {
|
277 |
"type": "content_block_delta",
|
278 |
"index": 0,
|
279 |
"delta": {"type": "text_delta", "text": f"Error during socket connection: {str(e)}"},
|
280 |
})
|
281 |
finally:
|
282 |
+
heartbeat_task.cancel()
|
283 |
if sio.connected:
|
284 |
+
await sio.disconnect()
|
285 |
|
286 |
+
yield send_event("content_block_stop", {"type": "content_block_stop", "index": 0})
|
287 |
+
yield send_event("message_delta", {
|
288 |
"type": "message_delta",
|
289 |
"delta": {"stop_reason": "end_turn", "stop_sequence": None},
|
290 |
"usage": {"output_tokens": total_output_tokens},
|
291 |
})
|
292 |
+
yield send_event("message_stop", {"type": "message_stop"})
|
293 |
|
294 |
return Response(generate(), content_type='text/event-stream')
|
295 |
|
|
|
413 |
})
|
414 |
if not API_KEY:
|
415 |
logger.warning("PPLX_KEY environment variable is not set", extra={'event_type': 'config_warning'})
|
416 |
+
app.run(host='0.0.0.0', port=port, debug=True)
|
417 |
|
418 |
|