bizyengine 1.2.2__tar.gz → 1.2.4__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (87) hide show
  1. {bizyengine-1.2.2 → bizyengine-1.2.4}/PKG-INFO +1 -1
  2. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/bizy_server/api_client.py +142 -0
  3. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/bizy_server/errno.py +15 -0
  4. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/bizy_server/server.py +159 -0
  5. bizyengine-1.2.4/bizyengine/bizy_server/stream_response.py +268 -0
  6. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/bizyair_extras/__init__.py +1 -0
  7. bizyengine-1.2.4/bizyengine/bizyair_extras/nodes_nunchaku.py +193 -0
  8. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/core/commands/servers/prompt_server.py +2 -2
  9. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/core/common/client.py +44 -31
  10. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/core/common/env_var.py +6 -3
  11. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/core/configs/models.json +1 -0
  12. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/core/configs/models.yaml +11 -0
  13. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/misc/auth.py +15 -26
  14. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/misc/utils.py +2 -2
  15. bizyengine-1.2.4/bizyengine/version.txt +1 -0
  16. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine.egg-info/PKG-INFO +1 -1
  17. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine.egg-info/SOURCES.txt +2 -0
  18. bizyengine-1.2.2/bizyengine/version.txt +0 -1
  19. {bizyengine-1.2.2 → bizyengine-1.2.4}/MANIFEST.in +0 -0
  20. {bizyengine-1.2.2 → bizyengine-1.2.4}/README.md +0 -0
  21. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/__init__.py +0 -0
  22. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/bizy_server/__init__.py +0 -0
  23. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/bizy_server/error_handler.py +0 -0
  24. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/bizy_server/execution.py +0 -0
  25. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/bizy_server/profile.py +0 -0
  26. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/bizy_server/resp.py +0 -0
  27. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/bizy_server/utils.py +0 -0
  28. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/bizyair_extras/nodes_advanced_refluxcontrol.py +0 -0
  29. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/bizyair_extras/nodes_cogview4.py +0 -0
  30. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/bizyair_extras/nodes_comfyui_detail_daemon.py +0 -0
  31. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/bizyair_extras/nodes_comfyui_instantid.py +0 -0
  32. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/bizyair_extras/nodes_comfyui_layerstyle_advance.py +0 -0
  33. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/bizyair_extras/nodes_comfyui_pulid_flux.py +0 -0
  34. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/bizyair_extras/nodes_controlnet.py +0 -0
  35. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/bizyair_extras/nodes_custom_sampler.py +0 -0
  36. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/bizyair_extras/nodes_dataset.py +0 -0
  37. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/bizyair_extras/nodes_differential_diffusion.py +0 -0
  38. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/bizyair_extras/nodes_flux.py +0 -0
  39. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/bizyair_extras/nodes_image_utils.py +0 -0
  40. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/bizyair_extras/nodes_ip2p.py +0 -0
  41. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/bizyair_extras/nodes_ipadapter_plus/__init__.py +0 -0
  42. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/bizyair_extras/nodes_ipadapter_plus/nodes_ipadapter_plus.py +0 -0
  43. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/bizyair_extras/nodes_janus_pro.py +0 -0
  44. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/bizyair_extras/nodes_kolors_mz/__init__.py +0 -0
  45. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/bizyair_extras/nodes_model_advanced.py +0 -0
  46. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/bizyair_extras/nodes_sd3.py +0 -0
  47. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/bizyair_extras/nodes_segment_anything.py +0 -0
  48. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/bizyair_extras/nodes_segment_anything_utils.py +0 -0
  49. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/bizyair_extras/nodes_testing_utils.py +0 -0
  50. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/bizyair_extras/nodes_trellis.py +0 -0
  51. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/bizyair_extras/nodes_ultimatesdupscale.py +0 -0
  52. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/bizyair_extras/nodes_upscale_model.py +0 -0
  53. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/bizyair_extras/nodes_wan_video.py +0 -0
  54. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/bizyair_extras/oauth_callback/main.py +0 -0
  55. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/bizyair_extras/route_bizyair_tools.py +0 -0
  56. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/core/__init__.py +0 -0
  57. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/core/commands/__init__.py +0 -0
  58. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/core/commands/base.py +0 -0
  59. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/core/commands/invoker.py +0 -0
  60. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/core/commands/processors/model_hosting_processor.py +0 -0
  61. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/core/commands/processors/prompt_processor.py +0 -0
  62. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/core/commands/servers/model_server.py +0 -0
  63. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/core/common/__init__.py +0 -0
  64. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/core/common/caching.py +0 -0
  65. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/core/common/utils.py +0 -0
  66. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/core/configs/conf.py +0 -0
  67. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/core/data_types.py +0 -0
  68. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/core/image_utils.py +0 -0
  69. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/core/nodes_base.py +0 -0
  70. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/core/nodes_io.py +0 -0
  71. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/core/path_utils/__init__.py +0 -0
  72. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/core/path_utils/path_manager.py +0 -0
  73. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/core/path_utils/utils.py +0 -0
  74. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/misc/__init__.py +0 -0
  75. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/misc/llm.py +0 -0
  76. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/misc/mzkolors.py +0 -0
  77. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/misc/nodes.py +0 -0
  78. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/misc/nodes_controlnet_aux.py +0 -0
  79. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/misc/nodes_controlnet_union_sdxl.py +0 -0
  80. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/misc/route_sam.py +0 -0
  81. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/misc/segment_anything.py +0 -0
  82. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine/misc/supernode.py +0 -0
  83. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine.egg-info/dependency_links.txt +0 -0
  84. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine.egg-info/requires.txt +0 -0
  85. {bizyengine-1.2.2 → bizyengine-1.2.4}/bizyengine.egg-info/top_level.txt +0 -0
  86. {bizyengine-1.2.2 → bizyengine-1.2.4}/pyproject.toml +0 -0
  87. {bizyengine-1.2.2 → bizyengine-1.2.4}/setup.cfg +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: bizyengine
3
- Version: 1.2.2
3
+ Version: 1.2.4
4
4
  Summary: [a/BizyAir](https://github.com/siliconflow/BizyAir) Comfy Nodes that can run in any environment.
5
5
  Author-email: SiliconFlow <yaochi@siliconflow.cn>
6
6
  Project-URL: Repository, https://github.com/siliconflow/BizyAir
@@ -1,3 +1,5 @@
1
+ import asyncio
2
+ import json
1
3
  import os
2
4
  import urllib
3
5
 
@@ -19,6 +21,10 @@ version_path = os.path.join(os.path.dirname(__file__), "..", "version.txt")
19
21
  with open(version_path, "r") as file:
20
22
  CLIENT_VERSION = file.read().strip()
21
23
 
24
+ # API请求地址
25
+ QWEN_MODEL_API_URL = "https://api.siliconflow.cn/v1/chat/completions"
26
+ QWEN_IMAGE_API_URL = "https://api.siliconflow.cn/v1/images/generations"
27
+
22
28
 
23
29
  class APIClient:
24
30
  def __init__(self):
@@ -1071,3 +1077,139 @@ class APIClient:
1071
1077
  except Exception as e:
1072
1078
  print(f"\033[31m[BizyAir]\033[0m Fail to get recent cost: {str(e)}")
1073
1079
  return None, errnos.GET_RECENT_COST
1080
+
1081
+ async def forward_model_request(self, request_data):
1082
+ try:
1083
+ import asyncio
1084
+ import json
1085
+
1086
+ from .stream_response import StreamResponse
1087
+
1088
+ request_data["stream"] = True
1089
+ # 参数检查
1090
+ if "messages" not in request_data:
1091
+ return None, errnos.MODEL_API_ERROR
1092
+
1093
+ if (
1094
+ not isinstance(request_data["messages"], list)
1095
+ or len(request_data["messages"]) == 0
1096
+ ):
1097
+ return None, errnos.MODEL_API_ERROR
1098
+
1099
+ # 获取用户API密钥
1100
+ api_key = get_api_key()
1101
+
1102
+ # 创建自定义流式响应对象,设置60秒超时
1103
+ response_stream = StreamResponse(request_data, timeout=60)
1104
+
1105
+ # 异步启动连接和数据请求
1106
+ asyncio.create_task(response_stream.connect_and_request(api_key))
1107
+
1108
+ # 返回流式响应对象
1109
+ return response_stream, None
1110
+
1111
+ except Exception as e:
1112
+ print(f"\033[31m[BizyAir]\033[0m Model API forwarding failed: {str(e)}")
1113
+ return None, errnos.MODEL_API_ERROR
1114
+
1115
+ async def stream_model_response(self, response):
1116
+ """
1117
+ 流式传输模型响应
1118
+
1119
+ Args:
1120
+ response: 从API获取的StreamResponse对象
1121
+
1122
+ Returns:
1123
+ 一个异步生成器,用于流式传输响应数据
1124
+ """
1125
+ # 生成请求ID用于日志
1126
+ req_id = f"stream-{id(response)}"
1127
+
1128
+ # 验证response对象是否有效
1129
+ if not response:
1130
+ print(f"\033[31m[BizyAir-{req_id}]\033[0m 无效的response对象 (为None)")
1131
+ yield b'data: {"error": "Invalid response object (None)"}\n\n'
1132
+ yield b"data: [DONE]\n\n"
1133
+ return
1134
+
1135
+ if not hasattr(response, "content") or not hasattr(
1136
+ response.content, "iter_any"
1137
+ ):
1138
+ print(
1139
+ f"\033[31m[BizyAir-{req_id}]\033[0m 无效的response对象 (缺少必要方法)"
1140
+ )
1141
+ yield b'data: {"error": "Invalid response object (missing methods)"}\n\n'
1142
+ yield b"data: [DONE]\n\n"
1143
+ return
1144
+
1145
+ error_occurred = False
1146
+ chunks_forwarded = 0
1147
+ chunk = None # 定义在外部,以便finally块可以访问
1148
+
1149
+ try:
1150
+ # 使用StreamResponse.iter_any方法创建流式传输生成器
1151
+ async for chunk in response.content.iter_any():
1152
+ chunks_forwarded += 1
1153
+
1154
+ yield chunk
1155
+ # 是否是结束标记
1156
+ if b"data: [DONE]" in chunk:
1157
+ break
1158
+ # 是否包含错误
1159
+ if b'data: {"error"' in chunk or b'data: {"message":' in chunk:
1160
+ error_occurred = True
1161
+
1162
+ except Exception as e:
1163
+ error_occurred = True
1164
+ error_msg = f"读取流式数据失败: {str(e)}"
1165
+ print(f"\033[31m[BizyAir-{req_id}]\033[0m {error_msg}")
1166
+
1167
+ # 返回错误消息
1168
+ error_json = json.dumps({"error": error_msg})
1169
+ yield f"data: {error_json}\n\n".encode("utf-8")
1170
+
1171
+ # 如果还没有发送过DONE,则发送
1172
+ if chunks_forwarded == 0 or (chunk and not b"data: [DONE]" in chunk):
1173
+ yield b"data: [DONE]\n\n"
1174
+
1175
+ finally:
1176
+ # 如果没有发生错误,但也没有转发任何数据,发送一个空响应
1177
+ if not error_occurred and chunks_forwarded == 0:
1178
+ yield b'data: {"content": ""}\n\n'
1179
+ yield b"data: [DONE]\n\n"
1180
+
1181
+ async def release(self):
1182
+ """关闭连接"""
1183
+ if self.writer and not self.closed:
1184
+ self.writer.close()
1185
+ try:
1186
+ await self.writer.wait_closed()
1187
+ except Exception:
1188
+ pass
1189
+ self.closed = True
1190
+
1191
+ async def forward_image_request(self, request_data):
1192
+ try:
1193
+ api_key = get_api_key()
1194
+ headers = {
1195
+ "Content-Type": "application/json",
1196
+ "Authorization": f"Bearer {api_key}",
1197
+ }
1198
+ # 创建异步HTTP会话
1199
+ async with aiohttp.ClientSession() as session:
1200
+ async with session.post(
1201
+ QWEN_IMAGE_API_URL, headers=headers, json=request_data
1202
+ ) as response:
1203
+ # 读取并解析响应
1204
+ if response.status != 200:
1205
+ error_text = await response.text()
1206
+ print(
1207
+ f"\033[31m[BizyAir]\033[0m Image generation failed: {error_text}"
1208
+ )
1209
+ return None, errnos.MODEL_API_ERROR
1210
+
1211
+ result = await response.json()
1212
+ return result, None
1213
+ except Exception as e:
1214
+ print(f"\033[31m[BizyAir]\033[0m Image generation request failed: {str(e)}")
1215
+ return None, errnos.MODEL_API_ERROR
@@ -442,3 +442,18 @@ class errnos:
442
442
  GET_RECENT_COST = ErrorNo(
443
443
  500, 500157, None, {"en": "Failed to get recent cost", "zh": "获取最近消费失败"}
444
444
  )
445
+
446
+ MODEL_API_ERROR = ErrorNo(
447
+ 500, 500158, None, {"en": "Failed to call model API", "zh": "调用模型API失败"}
448
+ )
449
+
450
+ MODEL_API_TIMEOUT = ErrorNo(
451
+ 500, 500159, None, {"en": "Model API request timeout", "zh": "模型API请求超时"}
452
+ )
453
+
454
+ STREAMING_CONNECTION_ERROR = ErrorNo(
455
+ 500,
456
+ 500160,
457
+ None,
458
+ {"en": "Streaming connection closed unexpectedly", "zh": "流式连接异常关闭"},
459
+ )
@@ -1,5 +1,6 @@
1
1
  import asyncio
2
2
  import configparser
3
+ import json
3
4
  import logging
4
5
  import os
5
6
  import shutil
@@ -23,6 +24,7 @@ COMMUNITY_API = f"{API_PREFIX}/community"
23
24
  MODEL_HOST_API = f"{API_PREFIX}/modelhost"
24
25
  USER_API = f"{API_PREFIX}/user"
25
26
  INVOICE_API = f"{API_PREFIX}/invoices"
27
+ MODEL_API = f"{API_PREFIX}/model"
26
28
 
27
29
  logging.basicConfig(level=logging.DEBUG)
28
30
 
@@ -957,6 +959,163 @@ class BizyAirServer:
957
959
 
958
960
  return OKResponse(resp)
959
961
 
962
+ @self.prompt_server.routes.post(f"/{MODEL_API}/chat")
963
+ async def chat_completions(request):
964
+ response = None # 确保变量在退出前定义
965
+ resp = None # 响应对象引用
966
+ req_id = f"req-{id(request)}" # 为请求生成唯一ID
967
+
968
+ try:
969
+ # 解析请求数据
970
+ request_data = await request.json()
971
+
972
+ # 转发请求到模型服务
973
+ response, err = await self.api_client.forward_model_request(
974
+ request_data
975
+ )
976
+ if err is not None:
977
+ print(
978
+ f"\033[31m[聊天请求-{req_id}]\033[0m 转发请求失败: {err.message}"
979
+ )
980
+ return ErrResponse(err)
981
+
982
+ # 创建并准备流式响应
983
+ resp = aiohttp.web.StreamResponse(
984
+ status=200,
985
+ reason="OK",
986
+ headers={
987
+ "Content-Type": "text/event-stream",
988
+ "Cache-Control": "no-cache",
989
+ "Connection": "keep-alive",
990
+ "X-Accel-Buffering": "no", # 禁用Nginx缓冲
991
+ },
992
+ )
993
+ await resp.prepare(request)
994
+
995
+ # 确保异步生成器正确创建
996
+ generator = self.api_client.stream_model_response(response)
997
+ if not generator:
998
+ print(f"\033[31m[聊天请求-{req_id}]\033[0m 创建异步生成器失败")
999
+ # 返回错误响应
1000
+ if not resp.prepared:
1001
+ return ErrResponse(errnos.MODEL_API_ERROR)
1002
+ else:
1003
+ # 如果响应已准备,发送错误数据帧
1004
+ try:
1005
+ error_msg = json.dumps({"error": "创建数据流失败"})
1006
+ await resp.write(f"data: {error_msg}\n\n".encode("utf-8"))
1007
+ await resp.write(b"data: [DONE]\n\n")
1008
+ except Exception as e:
1009
+ print(
1010
+ f"\033[31m[聊天请求-{req_id}]\033[0m 写入错误消息时出错: {str(e)}"
1011
+ )
1012
+ finally:
1013
+ try:
1014
+ await resp.write_eof()
1015
+ except:
1016
+ pass
1017
+ return resp
1018
+
1019
+ # 开始流式传输
1020
+ any_chunk_sent = False # 跟踪是否发送了任何数据块
1021
+
1022
+ try:
1023
+ async for chunk in generator:
1024
+ if chunk:
1025
+ try:
1026
+ await resp.write(chunk)
1027
+ any_chunk_sent = True
1028
+ await resp.drain() # 确保数据被立即发送
1029
+ except (
1030
+ ConnectionResetError,
1031
+ ConnectionError,
1032
+ aiohttp.ClientOSError,
1033
+ ) as e:
1034
+ print(
1035
+ f"\033[31m[聊天请求-{req_id}]\033[0m 流式传输中连接错误: {str(e)}"
1036
+ )
1037
+ break
1038
+ except Exception as e:
1039
+ print(f"\033[31m[聊天请求-{req_id}]\033[0m 流式传输错误: {str(e)}")
1040
+ # 如果尚未发送任何数据块,尝试发送错误信息
1041
+ if not any_chunk_sent and not resp.prepared:
1042
+ return ErrResponse(errnos.MODEL_API_ERROR)
1043
+ elif not any_chunk_sent:
1044
+ try:
1045
+ error_msg = json.dumps({"error": f"流式传输错误: {str(e)}"})
1046
+ await resp.write(f"data: {error_msg}\n\n".encode("utf-8"))
1047
+ await resp.write(b"data: [DONE]\n\n")
1048
+ except Exception as write_err:
1049
+ print(
1050
+ f"\033[31m[聊天请求-{req_id}]\033[0m 写入错误消息时出错: {str(write_err)}"
1051
+ )
1052
+
1053
+ # 检查是否发送了数据
1054
+ if not any_chunk_sent:
1055
+ print(
1056
+ f"\033[33m[聊天请求-{req_id}]\033[0m 警告: 没有发送任何数据块"
1057
+ )
1058
+ try:
1059
+ await resp.write(b'data: {"content": ""}\n\n')
1060
+ await resp.write(b"data: [DONE]\n\n")
1061
+ except Exception as e:
1062
+ print(
1063
+ f"\033[31m[聊天请求-{req_id}]\033[0m 写入空响应时出错: {str(e)}"
1064
+ )
1065
+
1066
+ try:
1067
+ await resp.write_eof()
1068
+ except Exception as e:
1069
+ print(
1070
+ f"\033[31m[聊天请求-{req_id}]\033[0m 结束响应时出错: {str(e)}"
1071
+ )
1072
+
1073
+ return resp
1074
+
1075
+ except Exception as e:
1076
+ print(
1077
+ f"\033[31m[聊天请求-{req_id}]\033[0m 处理请求时发生错误: {str(e)}"
1078
+ )
1079
+ # 如果响应已经准备好,尝试发送错误信息
1080
+ if resp and resp.prepared:
1081
+ try:
1082
+ error_msg = json.dumps({"error": f"服务器错误: {str(e)}"})
1083
+ await resp.write(f"data: {error_msg}\n\n".encode("utf-8"))
1084
+ await resp.write(b"data: [DONE]\n\n")
1085
+ await resp.write_eof()
1086
+ except:
1087
+ pass
1088
+ return resp
1089
+
1090
+ return ErrResponse(errnos.MODEL_API_ERROR)
1091
+
1092
+ finally:
1093
+ # 确保所有资源被释放
1094
+ if response and hasattr(response, "release"):
1095
+ try:
1096
+ await response.release()
1097
+ except Exception as e:
1098
+ print(
1099
+ f"\033[31m[聊天请求-{req_id}]\033[0m 关闭连接时出错: {str(e)}"
1100
+ )
1101
+
1102
+ @self.prompt_server.routes.post(f"/{MODEL_API}/images")
1103
+ async def image_generations(request):
1104
+ try:
1105
+ # 解析请求数据
1106
+ request_data = await request.json()
1107
+
1108
+ # 转发图像生成请求
1109
+ result, err = await self.api_client.forward_image_request(request_data)
1110
+ if err is not None:
1111
+ return ErrResponse(err)
1112
+
1113
+ # 返回结果
1114
+ return OKResponse(result)
1115
+
1116
+ except Exception:
1117
+ return ErrResponse(errnos.MODEL_API_ERROR)
1118
+
960
1119
  async def send_json(self, event, data, sid=None):
961
1120
  message = {"type": event, "data": data}
962
1121
 
@@ -0,0 +1,268 @@
1
+ import asyncio
2
+ import json
3
+ import ssl
4
+ from typing import Any, Dict, List, Optional
5
+
6
+
7
+ class ConnectionState:
8
+ """连接状态枚举"""
9
+
10
+ INIT = "初始化"
11
+ CONNECTING = "连接中"
12
+ CONNECTED = "已连接"
13
+ READING = "读取中"
14
+ CLOSING = "关闭中"
15
+ CLOSED = "已关闭"
16
+ ERROR = "错误"
17
+
18
+
19
+ class StreamResponse:
20
+ """自定义流式响应类,支持状态管理和错误处理"""
21
+
22
+ def __init__(self, request_data: Dict[str, Any], timeout: int = 60):
23
+ # 请求和连接数据
24
+ self.request_data = request_data
25
+ self.reader = None
26
+ self.writer = None
27
+ self.timeout = timeout
28
+
29
+ # 连接状态管理
30
+ self.state = ConnectionState.INIT
31
+ self.state_lock = asyncio.Lock() # 状态修改锁
32
+ self.connection_event = asyncio.Event()
33
+ self.error = None
34
+
35
+ # 数据相关
36
+ self.content = self
37
+ self.buffer = b""
38
+ self.chunks_received = 0
39
+ self.done_received = False
40
+
41
+ # 调试信息
42
+ self.debug_id = id(self) # 用于日志区分不同实例
43
+
44
+ async def _set_state(self, new_state: str, error: Exception = None) -> None:
45
+ """线程安全地设置连接状态"""
46
+ async with self.state_lock:
47
+ self.state = new_state
48
+ if error:
49
+ self.error = error
50
+
51
+ # 如果状态为已连接,触发连接事件
52
+ if new_state == ConnectionState.CONNECTED:
53
+ self.connection_event.set()
54
+
55
+ # 如果状态为错误,也触发连接事件,让等待连接的代码继续执行并处理错误
56
+ if new_state == ConnectionState.ERROR:
57
+ self.connection_event.set()
58
+
59
+ async def connect_and_request(self, api_key: str):
60
+ """建立连接并发送请求,包含严格的状态管理"""
61
+ # 设置状态为连接中
62
+ await self._set_state(ConnectionState.CONNECTING)
63
+
64
+ try:
65
+ # 建立SSL连接,添加超时控制
66
+ connect_coro = asyncio.open_connection("api.siliconflow.cn", 443, ssl=True)
67
+ self.reader, self.writer = await asyncio.wait_for(
68
+ connect_coro, timeout=self.timeout
69
+ )
70
+
71
+ # 准备HTTP请求
72
+ json_data = json.dumps(self.request_data)
73
+ request = (
74
+ f"POST /v1/chat/completions HTTP/1.1\r\n"
75
+ f"Host: api.siliconflow.cn\r\n"
76
+ f"Authorization: Bearer {api_key}\r\n"
77
+ f"Accept: text/event-stream\r\n"
78
+ f"Content-Type: application/json\r\n"
79
+ f"Content-Length: {len(json_data)}\r\n"
80
+ f"Connection: keep-alive\r\n"
81
+ f"\r\n"
82
+ f"{json_data}"
83
+ )
84
+
85
+ # 检查是否仍在连接状态
86
+ async with self.state_lock:
87
+ if self.state != ConnectionState.CONNECTING:
88
+ raise Exception(f"连接状态已变更为 {self.state},无法发送请求")
89
+
90
+ # 发送请求
91
+ self.writer.write(request.encode("utf-8"))
92
+ await asyncio.wait_for(self.writer.drain(), timeout=self.timeout)
93
+
94
+ # 读取HTTP响应头
95
+ response_line = await asyncio.wait_for(
96
+ self.reader.readline(), timeout=self.timeout
97
+ )
98
+ if not response_line:
99
+ raise Exception("服务器关闭了连接")
100
+
101
+ status_line = response_line.decode("utf-8").strip()
102
+
103
+ if not status_line.startswith("HTTP/1.1 200"):
104
+ raise Exception(f"API请求失败: {status_line}")
105
+
106
+ # 读取响应头
107
+ headers = {}
108
+ while True:
109
+ # 再次检查状态
110
+ async with self.state_lock:
111
+ if self.state != ConnectionState.CONNECTING:
112
+ raise Exception(
113
+ f"连接状态已变更为 {self.state},停止读取响应头"
114
+ )
115
+
116
+ line = await asyncio.wait_for(
117
+ self.reader.readline(), timeout=self.timeout
118
+ )
119
+ if line == b"\r\n" or not line:
120
+ break
121
+
122
+ try:
123
+ header_line = line.decode("utf-8").strip()
124
+ if ":" in header_line:
125
+ key, value = header_line.split(":", 1)
126
+ headers[key.strip()] = value.strip()
127
+ except Exception:
128
+ pass
129
+
130
+ # 成功连接
131
+ await self._set_state(ConnectionState.CONNECTED)
132
+
133
+ except asyncio.TimeoutError as e:
134
+ await self._set_state(ConnectionState.ERROR, e)
135
+ await self.release()
136
+ raise Exception(f"连接超时 (>{self.timeout}秒)") from e
137
+
138
+ except Exception as e:
139
+ await self._set_state(ConnectionState.ERROR, e)
140
+ await self.release()
141
+ raise Exception(f"连接失败: {str(e)}") from e
142
+
143
+ async def iter_any(self):
144
+ """模拟aiohttp的iter_any方法"""
145
+ # 等待连接完成
146
+ if not self.connection_event.is_set():
147
+ try:
148
+ await asyncio.wait_for(
149
+ self.connection_event.wait(), timeout=self.timeout
150
+ )
151
+ except asyncio.TimeoutError as e:
152
+ await self._set_state(ConnectionState.ERROR, e)
153
+ await self.release()
154
+ raise Exception(f"等待连接超时 (>{self.timeout}秒)") from e
155
+
156
+ # 检查连接状态
157
+ async with self.state_lock:
158
+ if self.state == ConnectionState.ERROR:
159
+ error_msg = f"连接发生错误: {str(self.error)}"
160
+ raise Exception(error_msg) from self.error
161
+
162
+ if self.state != ConnectionState.CONNECTED:
163
+ error_msg = f"连接状态为 {self.state},无法读取数据"
164
+ raise Exception(error_msg)
165
+
166
+ # 设置状态为读取中
167
+ self.state = ConnectionState.READING
168
+
169
+ # 读取数据
170
+ try:
171
+ max_empty_chunks = 3 # 连续空块计数,防止无限循环
172
+ empty_chunk_count = 0
173
+
174
+ while True:
175
+ # 检查状态
176
+ async with self.state_lock:
177
+ if self.state != ConnectionState.READING:
178
+ break
179
+
180
+ try:
181
+ # 读取一段数据,添加超时控制
182
+ chunk = await asyncio.wait_for(
183
+ self.reader.read(1024), timeout=self.timeout
184
+ )
185
+
186
+ # 检测连续空块,防止无限循环
187
+ if not chunk:
188
+ empty_chunk_count += 1
189
+ if empty_chunk_count >= max_empty_chunks:
190
+ break
191
+ # 给个短暂的休息,避免CPU过度使用
192
+ await asyncio.sleep(0.05)
193
+ continue
194
+ else:
195
+ empty_chunk_count = 0
196
+
197
+ # 更新计数并生成数据
198
+ self.chunks_received += 1
199
+ yield chunk
200
+
201
+ # 检查是否接收到了结束标记 [DONE]
202
+ if b"data: [DONE]" in chunk:
203
+ self.done_received = True
204
+ break
205
+
206
+ except asyncio.TimeoutError as e:
207
+ await self._set_state(ConnectionState.ERROR, e)
208
+ break
209
+
210
+ except (ConnectionError, OSError) as e:
211
+ await self._set_state(ConnectionState.ERROR, e)
212
+ break
213
+
214
+ except Exception as e:
215
+ await self._set_state(ConnectionState.ERROR, e)
216
+ error_msg = f"读取数据失败: {str(e)}"
217
+ raise
218
+ finally:
219
+ # 关闭连接
220
+ await self.release()
221
+
222
+ async def release(self):
223
+ """安全关闭连接并释放资源"""
224
+ async with self.state_lock:
225
+ # 检查是否已经在关闭或已关闭状态
226
+ if self.state in [ConnectionState.CLOSING, ConnectionState.CLOSED]:
227
+ return
228
+
229
+ # 设置状态为关闭中
230
+ self.state = ConnectionState.CLOSING
231
+
232
+ # 关闭写入端
233
+ if self.writer:
234
+ writer = self.writer
235
+ self.writer = None # 清除引用
236
+ self.reader = None # 清除引用
237
+
238
+ try:
239
+ # 尝试优雅关闭
240
+ if not writer.is_closing():
241
+ # 先确保发送所有数据
242
+ try:
243
+ await asyncio.wait_for(writer.drain(), timeout=1.0)
244
+ except (asyncio.TimeoutError, ConnectionError, OSError):
245
+ # 忽略drain错误
246
+ pass
247
+
248
+ # 写入EOF标记
249
+ try:
250
+ writer.write_eof()
251
+ except (OSError, RuntimeError):
252
+ # 忽略无法写入EOF的错误
253
+ pass
254
+
255
+ # 关闭连接
256
+ writer.close()
257
+
258
+ # 等待连接关闭完成
259
+ try:
260
+ await asyncio.wait_for(writer.wait_closed(), timeout=2.0)
261
+ except (asyncio.TimeoutError, ConnectionError, OSError, RuntimeError):
262
+ # 忽略等待关闭的错误
263
+ pass
264
+ except Exception:
265
+ pass
266
+
267
+ # 设置为已关闭状态
268
+ await self._set_state(ConnectionState.CLOSED)
@@ -15,6 +15,7 @@ from .nodes_ipadapter_plus.nodes_ipadapter_plus import *
15
15
  from .nodes_janus_pro import *
16
16
  from .nodes_kolors_mz import *
17
17
  from .nodes_model_advanced import *
18
+ from .nodes_nunchaku import *
18
19
  from .nodes_sd3 import *
19
20
  from .nodes_segment_anything import *
20
21
  from .nodes_testing_utils import *