129 lines
5.6 KiB
Python
129 lines
5.6 KiB
Python
def generate():
|
|
start_time = time.time()
|
|
|
|
yield f"event: step\ndata: {json.dumps({'step': 'context', 'content': f'已收集上下文: {data_summary}'}, ensure_ascii=False)}\n\n"
|
|
yield f"event: step\ndata: {json.dumps({'step': 'prompt', 'content': full_prompt_display}, ensure_ascii=False)}\n\n"
|
|
|
|
try:
|
|
resp = fake_post(
|
|
"http://127.0.0.1:4000/v1/chat/completions",
|
|
json={"model": "default", "messages": messages, "stream": True, "max_tokens": 1024},
|
|
stream=True,
|
|
timeout=120
|
|
)
|
|
resp.raise_for_status()
|
|
|
|
last_keepalive = time.time()
|
|
full_text = ""
|
|
# 滚动 buffer: 分隔符可能跨 chunk (如 "===REAS" + "ONING==="),
|
|
# 保留尾部最长分隔符前缀长度的字符, 不立即下发
|
|
DELIM_R = "===REASONING==="
|
|
DELIM_C = "===CONCLUSION==="
|
|
holdback = max(len(DELIM_R), len(DELIM_C)) - 1
|
|
pending = ""
|
|
reasoning_started = False
|
|
conclusion_started = False
|
|
|
|
def emit_token(t):
|
|
return f"event: token\ndata: {json.dumps({'text': t}, ensure_ascii=False)}\n\n"
|
|
|
|
for line in resp.iter_lines():
|
|
if not line:
|
|
continue
|
|
line = line.decode("utf-8") if isinstance(line, bytes) else line
|
|
|
|
now = time.time()
|
|
if now - last_keepalive > 15:
|
|
yield ": keepalive\n\n"
|
|
last_keepalive = now
|
|
|
|
if line.startswith("data: "):
|
|
chunk_data = line[6:]
|
|
if chunk_data.strip() == "[DONE]":
|
|
break
|
|
try:
|
|
chunk = json.loads(chunk_data)
|
|
delta = chunk.get("choices", [{}])[0].get("delta", {})
|
|
token_text = delta.get("content", "")
|
|
except json.JSONDecodeError:
|
|
continue
|
|
if not token_text:
|
|
continue
|
|
|
|
full_text += token_text
|
|
pending += token_text
|
|
|
|
# 循环剥离 pending 里确定安全的部分:
|
|
# 遇到完整分隔符 -> 发 section_start 事件, 丢弃分隔符;
|
|
# 遇到分隔符前缀后缀 -> 留在 pending 等待后续 chunk
|
|
while pending:
|
|
idx_r = pending.find(DELIM_R)
|
|
idx_c = pending.find(DELIM_C)
|
|
# 找最早出现的完整分隔符
|
|
idx = -1
|
|
delim = None
|
|
if idx_r != -1 and (idx_c == -1 or idx_r < idx_c):
|
|
idx, delim = idx_r, DELIM_R
|
|
elif idx_c != -1:
|
|
idx, delim = idx_c, DELIM_C
|
|
|
|
if idx != -1:
|
|
# 分隔符前的正文是安全的, 直接下发
|
|
if idx > 0:
|
|
if not reasoning_started:
|
|
reasoning_started = True
|
|
yield f"event: reasoning_start\ndata: {{}}\n\n"
|
|
yield emit_token(pending[:idx])
|
|
pending = pending[idx + len(delim):]
|
|
if delim == DELIM_R:
|
|
if not reasoning_started:
|
|
reasoning_started = True
|
|
yield f"event: reasoning_start\ndata: {{}}\n\n"
|
|
else:
|
|
if not reasoning_started:
|
|
# 防御: 模型跳过 REASONING 直接给 CONCLUSION
|
|
reasoning_started = True
|
|
yield f"event: reasoning_start\ndata: {{}}\n\n"
|
|
if not conclusion_started:
|
|
conclusion_started = True
|
|
yield f"event: conclusion_start\ndata: {{}}\n\n"
|
|
continue
|
|
|
|
# 无完整分隔符: 检查尾部是否为某分隔符的前缀
|
|
keep = 0
|
|
tail = pending[-holdback:] if len(pending) > holdback else pending
|
|
for d in (DELIM_R, DELIM_C):
|
|
for k in range(min(len(tail), len(d) - 1), 0, -1):
|
|
if tail.endswith(d[:k]):
|
|
keep = max(keep, k)
|
|
break
|
|
safe = pending[:len(pending) - keep] if keep else pending
|
|
pending = pending[len(safe):]
|
|
if not safe:
|
|
break
|
|
if not reasoning_started:
|
|
reasoning_started = True
|
|
yield f"event: reasoning_start\ndata: {{}}\n\n"
|
|
yield emit_token(safe)
|
|
|
|
# 流结束: 残余 pending 全部下发
|
|
if pending:
|
|
if not reasoning_started:
|
|
reasoning_started = True
|
|
yield f"event: reasoning_start\ndata: {{}}\n\n"
|
|
yield emit_token(pending)
|
|
pending = ""
|
|
|
|
except Exception as e:
|
|
yield f"event: error\ndata: {json.dumps({'error': f'LLM call failed: {str(e)}'}, ensure_ascii=False)}\n\n"
|
|
return
|
|
|
|
# 防御: 模型没给分隔符时补发 section_start, 保证前端 tab 一定有入口
|
|
if not reasoning_started:
|
|
yield f"event: reasoning_start\ndata: {{}}\n\n"
|
|
if not conclusion_started:
|
|
yield f"event: conclusion_start\ndata: {{}}\n\n"
|
|
|
|
elapsed = round(time.time() - start_time, 1)
|
|
yield f"event: step\ndata: {json.dumps({'step': 'complete', 'content': f'分析完成, 总耗时 {elapsed} 秒'}, ensure_ascii=False)}\n\n"
|