letterm commited on
Commit
2acc97a
·
verified ·
1 Parent(s): f8f4a71

Upload 2 files

Browse files
Files changed (2) hide show
  1. Dockerfile +12 -0
  2. app.py +561 -0
Dockerfile ADDED
@@ -0,0 +1,12 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ FROM python:3.10-slim
2
+
3
+ WORKDIR /app
4
+
5
+ RUN pip install --no-cache-dir protobuf asyncio aiohttp websockets
6
+
7
+ COPY . .
8
+
9
+ ENV PORT=5200
10
+ EXPOSE 5200
11
+
12
+ CMD ["python", "app.py"]
app.py ADDED
@@ -0,0 +1,561 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Generated from trimmed zed.proto
2
+ from google.protobuf import descriptor as _descriptor
3
+ from google.protobuf import descriptor_pool as _descriptor_pool
4
+ from google.protobuf import runtime_version as _runtime_version
5
+ from google.protobuf import symbol_database as _symbol_database
6
+ from google.protobuf.internal import builder as _builder
7
+ _runtime_version.ValidateProtobufRuntimeVersion(
8
+ _runtime_version.Domain.PUBLIC,
9
+ 5,
10
+ 29,
11
+ 0,
12
+ '',
13
+ 'zed.proto'
14
+ )
15
+ # @@protoc_insertion_point(imports)
16
+
17
+ _sym_db = _symbol_database.Default()
18
+
19
+ DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\tzed.proto\x12\x0czed.messages\"&\n\x06PeerId\x12\x10\n\x08owner_id\x18\x01 \x01(\r\x12\n\n\x02id\x18\x02 \x01(\r\"\xe2\x05\n\x08\x45nvelope\x12\n\n\x02id\x18\x01 \x01(\r\x12\x1a\n\rresponding_to\x18\x02 \x01(\rH\x01\x88\x01\x01\x12\x35\n\x12original_sender_id\x18\x03 \x01(\x0b\x32\x14.zed.messages.PeerIdH\x02\x88\x01\x01\x12\x14\n\x06\x61\x63k_id\x18\x8a\x02 \x01(\rH\x03\x88\x01\x01\x12$\n\x05hello\x18\x04 \x01(\x0b\x32\x13.zed.messages.HelloH\x00\x12 \n\x03\x61\x63k\x18\x05 \x01(\x0b\x32\x11.zed.messages.AckH\x00\x12$\n\x05\x65rror\x18\x06 \x01(\x0b\x32\x13.zed.messages.ErrorH\x00\x12\"\n\x04ping\x18\x07 \x01(\x0b\x32\x12.zed.messages.PingH\x00\x12\x46\n\x17\x61\x63\x63\x65pt_terms_of_service\x18\xef\x01 \x01(\x0b\x32".zed.messages.AcceptTermsOfServiceH\x00\x12W\n accept_terms_of_service_response\x18\xf0\x01 \x01(\x0b\x32*.zed.messages.AcceptTermsOfServiceResponseH\x00\x12\x33\n\rget_llm_token\x18\xeb\x01 \x01(\x0b\x32\x19.zed.messages.GetLlmTokenH\x00\x12\x44\n\x16get_llm_token_response\x18\xec\x01 \x01(\x0b\x32!.zed.messages.GetLlmTokenResponseH\x00\x42\t\n\x07payloadB\x10\n\x0e_responding_toB\x15\n\x13_original_sender_idB\t\n\x07_ack_idJ\x04\x08W\x10YJ\x06\x08\x9e\x01\x10\xa2\x01J\x06\x08\xa4\x01\x10\xa5\x01J\x06\x08\xa6\x01\x10\xaa\x01J\x06\x08\xb1\x01\x10\xba\x01J\x06\x08\xbc\x01\x10\xbd\x01J\x06\x08\xc1\x01\x10\xc4\x01J\x06\x08\xc5\x01\x10\xc6\x01J\x06\x08\xc8\x01\x10\xcb\x01J\x06\x08\xcd\x01\x10\xcf\x01J\x06\x08\xdd\x01\x10\xde\x01J\x06\x08\xe0\x01\x10\xe6\x01J\x06\x08\xf6\x01\x10\xf7\x01J\x06\x08\xf7\x01\x10\xff\x01J\x06\x08\xff\x01\x10\x81\x02\".\n\x05Hello\x12%\n\x07peer_id\x18\x01 \x01(\x0b\x32\x14.zed.messages.PeerId\"\x06\n\x04Ping\"\x05\n\x03\x41\x63k\"M\n\x05\x45rror\x12\x0f\n\x07message\x18\x01 \x01(\t\x12%\n\x04\x63ode\x18\x02 \x01(\x0e\x32\x17.zed.messages.ErrorCode\x12\x0c\n\x04tags\x18\x03 \x03(\t\"\x16\n\x14\x41\x63\x63\x65ptTermsOfService\"7\n\x1c\x41\x63\x63\x65ptTermsOfServiceResponse\x12\x17\n\x0f\x61\x63\x63\x65pted_tos_at\x18\x01 \x01(\x04\"\r\n\x0bGetLlmToken\"$\n\x13GetLlmTokenResponse\x12\r\n\x05token\x18\x01 \x01(\t*\xef\x02\n\tErrorCode\x12\x0c\n\x08Internal\x10\x00\x12\x11\n\rNoSuchChannel\x10\x01\x12\x10\n\x0c\x44isconnected\x10\x02\x12\r\n\tSignedOut\x10\x03\x12\x13\n\x0fUpgradeRequired\x10\x04\x12\r\n\tForbidden\x10\x05\x12\x0c\n\x08NeedsCla\x10\x07\x12\x13\n\x0fNotARootChannel\x10\x08\x12\x14\n\x10\x42\x61\x64PublicNesting\x10\t\x12\x13\n\x0f\x43ircularNesting\x10\n\x12\x13\n\x0fWrongMoveTarget\x10\x0b\x12\x10\n\x0cUnsharedItem\x10\x0c\x12\x11\n\rNoSuchProject\x10\r\x12$\n DevServerProjectPathDoesNotExist\x10\x10\x12\x19\n\x15RemoteUpgradeRequired\x10\x11\x12\x15\n\x11RateLimitExceeded\x10\x12\x12\x10\n\x0c\x43ommitFailed\x10\x13\"\x04\x08\x06\x10\x06\"\x04\x08\x0e\x10\x0f\x62\x06proto3')
20
+
21
+ _globals = globals()
22
+ _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
23
+ _builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'zed_pb2', _globals)
24
+ if not _descriptor._USE_C_DESCRIPTORS:
25
+ DESCRIPTOR._loaded_options = None
26
+ _globals['_ERRORCODE']._serialized_start=1085
27
+ _globals['_ERRORCODE']._serialized_end=1452
28
+ _globals['_PEERID']._serialized_start=27
29
+ _globals['_PEERID']._serialized_end=65
30
+ _globals['_ENVELOPE']._serialized_start=68
31
+ _globals['_ENVELOPE']._serialized_end=806
32
+ _globals['_HELLO']._serialized_start=808
33
+ _globals['_HELLO']._serialized_end=854
34
+ _globals['_PING']._serialized_start=856
35
+ _globals['_PING']._serialized_end=862
36
+ _globals['_ACK']._serialized_start=864
37
+ _globals['_ACK']._serialized_end=869
38
+ _globals['_ERROR']._serialized_start=871
39
+ _globals['_ERROR']._serialized_end=948
40
+ _globals['_ACCEPTTERMSOFSERVICE']._serialized_start=950
41
+ _globals['_ACCEPTTERMSOFSERVICE']._serialized_end=972
42
+ _globals['_ACCEPTTERMSOFSERVICERESPONSE']._serialized_start=974
43
+ _globals['_ACCEPTTERMSOFSERVICERESPONSE']._serialized_end=1029
44
+ _globals['_GETLLMTOKEN']._serialized_start=1031
45
+ _globals['_GETLLMTOKEN']._serialized_end=1044
46
+ _globals['_GETLLMTOKENRESPONSE']._serialized_start=1046
47
+ _globals['_GETLLMTOKENRESPONSE']._serialized_end=1082
48
+
49
+ # Start of the actual script
50
+ import os
51
+ import json
52
+ import ssl
53
+ import time
54
+ import asyncio
55
+ import logging
56
+ import aiohttp
57
+ from aiohttp import web
58
+ import zstandard as zstd
59
+ from websockets.asyncio.client import connect
60
+ from websockets.exceptions import ConnectionClosed
61
+ import uuid
62
+
63
+ from google.protobuf.json_format import MessageToDict
64
+
65
+ Envelope = _sym_db.GetSymbol('zed.messages.Envelope')
66
+
67
+ logging.basicConfig(
68
+ level=logging.INFO,
69
+ format='%(levelname)s: %(message)s'
70
+ )
71
+ logger = logging.getLogger(__name__)
72
+
73
+
74
+
75
+ CONFIG = {
76
+ "API":{
77
+ "BASE_URL": "https://zed.dev",
78
+ "API_KEY": os.getenv("API_KEY","sk-123456"),
79
+ "BASE_API_URL": "https://collab.zed.dev",
80
+ "WS_URL": "wss://collab.zed.dev/rpc",
81
+ "LLM_API_URL": "https://llm.zed.dev/completion",
82
+ },
83
+ "LOGIN":{
84
+ "USER_ID": os.getenv("ZED_USER_ID"),
85
+ "AUTH": os.getenv("ZED_AUTH_TOKEN")
86
+ },
87
+ "SERVER":{
88
+ "PORT": os.getenv("PORT",5200),
89
+ "TOKEN_EXPIRY_WARNING_MINUTES": 50
90
+ },
91
+ "MODELS":{
92
+ "claude-3-5-sonnet-20241022":"claude-3-5-sonnet-latest",
93
+ "claude-3-7-sonnet-20250219":"claude-3-7-sonnet-20250219"
94
+ }
95
+ }
96
+
97
+ highest_message_id = 0
98
+ llm_token = None
99
+ token_timestamp = None
100
+ server_peer_id = None
101
+ active_websocket = None
102
+ proxy_server_running = False
103
+
104
+ class MessageProcessor:
105
+ @staticmethod
106
+ def create_chat_response(message, model, is_stream=False):
107
+ base_response = {
108
+ "id": f"chatcmpl-{uuid.uuid4()}",
109
+ "created": int(time.time()),
110
+ "model": model
111
+ }
112
+
113
+ if is_stream:
114
+ return {
115
+ **base_response,
116
+ "object": "chat.completion.chunk",
117
+ "choices": [{
118
+ "index": 0,
119
+ "delta": {
120
+ "content": message
121
+ }
122
+ }]
123
+ }
124
+
125
+ return {
126
+ **base_response,
127
+ "object": "chat.completion",
128
+ "choices": [{
129
+ "index": 0,
130
+ "message": {
131
+ "role": "assistant",
132
+ "content": message
133
+ },
134
+ "finish_reason": "stop"
135
+ }],
136
+ "usage": None
137
+ }
138
+
139
+ def decode_envelope(data):
140
+ try:
141
+ dctx = zstd.ZstdDecompressor()
142
+ decompressed_data = b''
143
+ with dctx.stream_reader(data) as reader:
144
+ while True:
145
+ chunk = reader.read(8192)
146
+ if not chunk:
147
+ break
148
+ decompressed_data += chunk
149
+ envelope = Envelope()
150
+ envelope.ParseFromString(decompressed_data)
151
+ return MessageToDict(envelope, preserving_proto_field_name=True)
152
+ except Exception as e:
153
+ hex_preview = ' '.join(f'{byte:02x}' for byte in data[:20]) + ('...' if len(data) > 20 else '')
154
+ logger.error(f"无法解码消息: {e}; 数据预览: {hex_preview}")
155
+ return {"error": f"无法解码消息: {e}"}
156
+
157
+ def compress_protobuf(data):
158
+ return zstd.ZstdCompressor(level=-7).compress(data)
159
+
160
+ def create_message(message_type):
161
+ global highest_message_id
162
+ highest_message_id += 1
163
+ message_id = highest_message_id
164
+ envelope = Envelope(id=highest_message_id)
165
+
166
+ getattr(envelope, message_type).SetInParent()
167
+ return compress_protobuf(envelope.SerializeToString()), message_id
168
+
169
+ async def ping_periodically(websocket):
170
+ while True:
171
+ try:
172
+ await websocket.ping()
173
+ await asyncio.sleep(1)
174
+ except Exception as e:
175
+ logger.error(f"发送ping错误: {e}")
176
+ break
177
+
178
+ async def handle_messages(websocket):
179
+ global server_peer_id, llm_token, token_timestamp, active_websocket
180
+ active_websocket = websocket
181
+ try:
182
+ async for message in websocket:
183
+ message_bytes = message.encode('utf-8') if isinstance(message, str) else message
184
+ decoded = decode_envelope(message_bytes)
185
+ if "hello" in decoded:
186
+ server_peer_id = decoded.get('hello', {}).get('peer_id')
187
+ elif "accept_terms_of_service_response" in decoded:
188
+ await request_llm_token(websocket)
189
+ elif ("get_llm_token_response" in decoded and
190
+ 'token' in decoded.get('get_llm_token_response', {})):
191
+ llm_token = decoded['get_llm_token_response']['token']
192
+ token_timestamp = time.time()
193
+ logger.info(f"LLM令牌收到 {time.ctime(token_timestamp)}")
194
+ if not proxy_server_running:
195
+ asyncio.create_task(start_proxy_server())
196
+ asyncio.create_task(monitor_token_expiration())
197
+ logger.info("关闭WebSocket连接,直到需要刷新令牌")
198
+ await websocket.close()
199
+ active_websocket = None
200
+ return
201
+ except ConnectionClosed:
202
+ logger.info("连接已关闭")
203
+ active_websocket = None
204
+
205
+ async def request_llm_token(websocket):
206
+ message, _ = create_message('get_llm_token')
207
+ logger.info("请求LLM令牌")
208
+ await websocket.send(message)
209
+
210
+ async def request_accept_terms_of_service(websocket):
211
+ message, _ = create_message('accept_terms_of_service')
212
+ logger.info("发送同意Zed服务条款")
213
+ await websocket.send(message)
214
+
215
+ def format_content(content):
216
+ if isinstance(content, str):
217
+ return [{"type": "text", "text": content}]
218
+ return content
219
+
220
+
221
+
222
+ async def process_message_content(content):
223
+ """
224
+ 处理消息内容,将不同类型的内容转换为字符串
225
+ """
226
+ if isinstance(content, str):
227
+ return content
228
+
229
+ if isinstance(content, list):
230
+ return '\n'.join([item.get('text', '') for item in content])
231
+
232
+ if isinstance(content, dict):
233
+ return content.get('text', None)
234
+
235
+ return None
236
+
237
+ async def transform_messages(request):
238
+ """
239
+ 转换消息格式,合并系统消息并处理消息结构
240
+ """
241
+ system_message = '' # 存储系统消息的变量
242
+ is_collecting_system_message = False # 是否正在收集系统消息
243
+ has_processed_system_messages = False # 是否已处理初始系统消息
244
+
245
+ converted_messages = []
246
+
247
+ for current in request.get('messages', []):
248
+ role = current.get('role')
249
+ current_content = await process_message_content(current.get('content'))
250
+
251
+ if current_content is None:
252
+ converted_messages.append(current)
253
+ continue
254
+
255
+ if role == 'system' and not has_processed_system_messages:
256
+ if not is_collecting_system_message:
257
+ # 第一次遇到system,开启收集
258
+ system_message = current_content
259
+ is_collecting_system_message = True
260
+ else:
261
+ # 继续遇到system,合并system消息
262
+ system_message += '\n' + current_content
263
+ continue
264
+
265
+ # 遇到非system消息
266
+ if is_collecting_system_message:
267
+ # 结束系统消息收集
268
+ is_collecting_system_message = False
269
+ has_processed_system_messages = True
270
+
271
+ # 如果已处理初始消息序列且再次遇到system,则转换role为user
272
+ if has_processed_system_messages and role == 'system':
273
+ role = 'user'
274
+
275
+ # 检查是否可以合并消息
276
+ if converted_messages and converted_messages[-1].get('role') == role:
277
+ converted_messages[-1]['content'][0]['text'] += '\r\n' + current_content
278
+ else:
279
+ converted_messages.append({
280
+ 'role': role,
281
+ 'content': [{'type': 'text', 'text': current_content}]
282
+ })
283
+
284
+ return {
285
+ 'messages': converted_messages,
286
+ 'system': system_message,
287
+ 'model': CONFIG['MODELS'].get(request.get('model'), "claude-3-5-sonnet-latest"),
288
+ 'max_tokens': request.get('max_tokens',8192),
289
+ 'temperature': max(0, min(request.get('temperature', 0), 1)),
290
+ 'top_p': max(0, min(request.get('top_p', 1), 1)),
291
+ 'top_k': max(0, min(request.get('top_k', 0), 500)),
292
+ 'stream': True
293
+ }
294
+
295
+ @web.middleware
296
+ async def auth_middleware(request, handler):
297
+ if CONFIG['API']['API_KEY']:
298
+ auth_header = request.headers.get('Authorization')
299
+ xapi_key_header = request.headers.get('x-api-key')
300
+
301
+ auth_password = None
302
+ if auth_header and auth_header.startswith('Bearer '):
303
+ auth_password = auth_header[7:]
304
+
305
+ if auth_password == CONFIG['API']['API_KEY'] or xapi_key_header == CONFIG['API']['API_KEY']:
306
+ return await handler(request)
307
+ else:
308
+ return web.json_response(
309
+ {"error": "Unauthorized"},
310
+ status=401
311
+ )
312
+
313
+ return await handler(request)
314
+
315
+ async def handle_models_request(request):
316
+ return web.json_response({
317
+ "object": "list",
318
+ "data": [
319
+ {
320
+ "id": model,
321
+ "object": "model",
322
+ "created": int(time.time()),
323
+ "owned_by": "zed"
324
+ }
325
+ for model in CONFIG["MODELS"].keys()
326
+ ]
327
+ })
328
+
329
+ async def handle_message_request(request):
330
+ global llm_token
331
+ if not llm_token:
332
+ return web.json_response({"error": "LLM令牌不可用"}, status=500)
333
+ try:
334
+ body = await request.json()
335
+ isClaudeAI = False
336
+ if request.path == '/v1/messages':
337
+ isClaudeAI = True
338
+ if "messages" in body:
339
+ for msg in body["messages"]:
340
+ if "content" in msg:
341
+ msg["content"] = format_content(msg["content"])
342
+ if "system" in body:
343
+ if isinstance(body["system"], list):
344
+ body["system"] = "\n".join([item["text"] for item in body["system"]])
345
+ if "model" in body:
346
+ body["model"] = CONFIG['MODELS'].get(body["model"], "claude-3-5-sonnet-latest")
347
+ else:
348
+ body = await transform_messages(body)
349
+ with open('request_payload222.json', 'w', encoding='utf-8') as f:
350
+ json.dump(body, f, ensure_ascii=False, indent=2)
351
+ headers = {"Content-Type": "application/json", "Authorization": f"Bearer {llm_token}"}
352
+ with open('ceshi.txt', 'w', encoding='utf-8') as f:
353
+ f.write(llm_token + '\n')
354
+ payload = {
355
+ "provider": "anthropic",
356
+ "model": body.get("model", "claude-3-5-sonnet-latest"),
357
+ "provider_request": body
358
+ }
359
+ # with open('ceshi.txt', 'w', encoding='utf-8') as f:
360
+ # f.write(json.dumps(body,ensure_ascii=False) + '\n')
361
+ if body.get("stream", False):
362
+ return await handle_streaming_request(request, headers, payload, isClaudeAI)
363
+ else:
364
+ return await handle_non_streaming_request(headers, payload, isClaudeAI)
365
+ except Exception as e:
366
+ logger.error(f"处理请求时发生错误: {e}")
367
+ return web.json_response({"error": str(e)}, status=500)
368
+
369
+ async def handle_non_streaming_request(headers, payload, isClaudeAI=False):
370
+ async with aiohttp.ClientSession() as session:
371
+ async with session.post(CONFIG['API']['LLM_API_URL'], headers=headers, json=payload) as r:
372
+ if r.status != 200:
373
+ text = await r.text()
374
+ logger.error(f"LLM API错误: {text}")
375
+ return web.json_response({"error": text}, status=r.status)
376
+ full_content, message_data = "", {}
377
+ async for line in r.content:
378
+ if not line:
379
+ continue
380
+ try:
381
+ event = json.loads(line.decode('utf-8').strip())
382
+ et = event.get('type')
383
+ if et == "message_start":
384
+ message_data = event.get('message', {})
385
+ elif et == "content_block_delta" and event.get('delta', {}).get('type') == "text_delta":
386
+ full_content += event['delta'].get('text', '')
387
+ elif et == "message_delta" and 'usage' in event:
388
+ message_data['usage'] = event.get('usage')
389
+ elif et == "message_stop":
390
+ break
391
+ except Exception as e:
392
+ logger.error(f"Error processing line: {e}")
393
+ if isClaudeAI:
394
+ message_data['content'] = [{"type": "text", "text": full_content}]
395
+ else:
396
+ message_data = MessageProcessor.create_chat_response(full_content, payload.get("model"), False)
397
+ return web.json_response(message_data)
398
+
399
+ async def handle_streaming_request(request, headers, payload, isClaudeAI=False):
400
+ response = web.StreamResponse()
401
+ response.headers['Content-Type'] = 'text/event-stream'
402
+ response.headers['Cache-Control'] = 'no-cache'
403
+ response.headers['Connection'] = 'keep-alive'
404
+ await response.prepare(request)
405
+ logger.info(f"开始处理流请求")
406
+ async with aiohttp.ClientSession() as session:
407
+ async with session.post(CONFIG['API']['LLM_API_URL'], headers=headers, json=payload) as api_response:
408
+ if api_response.status != 200:
409
+ error_text = await api_response.text()
410
+ logger.error(f"LLM API (stream)错误: {error_text}")
411
+ await response.write(f"data: {json.dumps({'error': error_text})}\n\n".encode())
412
+ await response.write(b"data: [DONE]\n\n")
413
+ return response
414
+ async for line in api_response.content:
415
+ try:
416
+ if line:
417
+ if isClaudeAI:
418
+ await response.write(f"data: {line.decode('utf-8')}\n\n".encode())
419
+ else:
420
+ try:
421
+ data = json.loads(line.decode('utf-8').strip())
422
+ if data.get('type') == "content_block_delta" and data.get('delta', {}).get('type') == "text_delta":
423
+ text = data['delta'].get('text', '')
424
+ message = MessageProcessor.create_chat_response(text, payload.get("model"), True)
425
+ await response.write(f"data: {json.dumps(message)}\n\n".encode())
426
+ except Exception as e:
427
+ logger.error(f"Error processing line: {e}")
428
+ except Exception as e:
429
+ logger.error(f"Error processing line: {e}")
430
+ await response.write(b"data: [DONE]\n\n")
431
+ return response
432
+
433
+
434
+ async def start_proxy_server():
435
+ global proxy_server_running
436
+ if proxy_server_running:
437
+ logger.info("代理服务器已运行,跳过启动")
438
+ return
439
+
440
+ proxy_server_running = True
441
+ app = web.Application(middlewares=[auth_middleware])
442
+ app.router.add_post('/v1/messages', handle_message_request)
443
+ app.router.add_post('/v1/chat/completions', handle_message_request)
444
+ app.router.add_get('/v1/models', handle_models_request)
445
+
446
+ async def health_check():
447
+ return web.json_response({
448
+ "status": "ok",
449
+ "message": "Zed LLM proxy is running"
450
+ })
451
+
452
+ app.router.add_get('/', health_check)
453
+
454
+ runner = web.AppRunner(app)
455
+ await runner.setup()
456
+ site = web.TCPSite(runner, 'localhost', CONFIG['SERVER']['PORT'])
457
+ await site.start()
458
+ logger.info(f"代理服务器启动 http://localhost:{CONFIG['SERVER']['PORT']}")
459
+ while True:
460
+ await asyncio.sleep(3600)
461
+
462
+ def is_token_expiring():
463
+ if not token_timestamp:
464
+ return False
465
+ return (time.time() - token_timestamp) / 60 >= CONFIG['SERVER']['TOKEN_EXPIRY_WARNING_MINUTES']
466
+
467
+ async def monitor_token_expiration():
468
+ while True:
469
+ await asyncio.sleep(60)
470
+ if is_token_expiring():
471
+ elapsed = int((time.time() - token_timestamp) / 60)
472
+ logger.warning(f"LLM令牌接近过期 (收到 {elapsed} 分钟前)")
473
+ if active_websocket is None:
474
+ logger.info("重新连接WebSocket以刷新令牌")
475
+ asyncio.create_task(reconnect_for_token_refresh())
476
+ return
477
+
478
+ async def reconnect_for_token_refresh():
479
+ try:
480
+ if not CONFIG['LOGIN']['USER_ID'] or not CONFIG['LOGIN']['AUTH']:
481
+ logger.error("用户ID或授权令牌未设置")
482
+ return
483
+ headers = {
484
+ "authorization": f"{CONFIG['LOGIN']['USER_ID']} {CONFIG['LOGIN']['AUTH']}",
485
+ "x-zed-protocol-version": "68",
486
+ "x-zed-app-version": "0.178.0",
487
+ "x-zed-release-channel": "stable"
488
+ }
489
+ ssl_context = ssl.create_default_context()
490
+ ssl_context.check_hostname = False
491
+ ssl_context.verify_mode = ssl.CERT_NONE
492
+
493
+ async for websocket in connect(CONFIG['API']['WS_URL'], additional_headers=headers, ssl=ssl_context):
494
+ try:
495
+ ping_task = asyncio.create_task(ping_periodically(websocket))
496
+ await asyncio.sleep(2)
497
+ await request_accept_terms_of_service(websocket)
498
+ await handle_messages(websocket)
499
+ break
500
+ except ConnectionClosed:
501
+ continue
502
+ except Exception as e:
503
+ logger.error(f"令牌刷新期间发生错误: {e}")
504
+ await asyncio.sleep(1)
505
+ continue
506
+ finally:
507
+ ping_task.cancel()
508
+ try:
509
+ await ping_task
510
+ except asyncio.CancelledError:
511
+ pass
512
+ except Exception as e:
513
+ logger.error(f"令牌刷新失败: {e}")
514
+
515
+ async def async_main():
516
+ if not CONFIG['LOGIN']['USER_ID'] or not CONFIG['LOGIN']['AUTH']:
517
+ logger.error("用户ID或授权令牌未设置")
518
+ return
519
+ headers = {
520
+ "authorization": f"{CONFIG['LOGIN']['USER_ID']} {CONFIG['LOGIN']['AUTH']}",
521
+ "x-zed-protocol-version": "68",
522
+ "x-zed-app-version": "0.178.0",
523
+ "x-zed-release-channel": "stable"
524
+ }
525
+ ssl_context = ssl.create_default_context()
526
+ ssl_context.check_hostname = False
527
+ ssl_context.verify_mode = ssl.CERT_NONE
528
+ logger.info("连接到Websocket服务器")
529
+ async for websocket in connect(CONFIG['API']['WS_URL'], additional_headers=headers, ssl=ssl_context):
530
+ try:
531
+ ping_task = asyncio.create_task(ping_periodically(websocket))
532
+ token_request_task = asyncio.create_task(delayed_token_request(websocket, delay=2))
533
+ await handle_messages(websocket)
534
+ break
535
+ except ConnectionClosed:
536
+ continue
537
+ except Exception as e:
538
+ logger.error(f"意外错误: {e}")
539
+ await asyncio.sleep(1)
540
+ continue
541
+ finally:
542
+ ping_task.cancel()
543
+ try:
544
+ await ping_task
545
+ except asyncio.CancelledError:
546
+ pass
547
+ token_request_task.cancel()
548
+ try:
549
+ await token_request_task
550
+ except asyncio.CancelledError:
551
+ pass
552
+
553
+ while True:
554
+ await asyncio.sleep(3600)
555
+
556
+ async def delayed_token_request(websocket, delay=2):
557
+ await asyncio.sleep(delay)
558
+ await request_accept_terms_of_service(websocket)
559
+
560
+ if __name__ == "__main__":
561
+ asyncio.run(async_main())