Merge pull request #3125 from xinnan-tech/py-fix-llm-chunk

fix: 修复频繁打断下生成器不会清理
This commit is contained in:
wengzh
2026-04-22 11:30:35 +08:00
committed by GitHub
4 changed files with 131 additions and 117 deletions
@@ -39,13 +39,12 @@ class LLMProvider(LLMProviderBase):
} }
# 发起 POST 请求 # 发起 POST 请求
response = requests.post(self.api_url, json=payload, headers=headers) with requests.post(self.api_url, json=payload, headers=headers) as response:
# 检查请求是否成功
response.raise_for_status()
# 检查请求是否成功 # 解析返回数据
response.raise_for_status() data = response.json()
# 解析返回数据
data = response.json()
speech = ( speech = (
data.get("response", {}) data.get("response", {})
.get("speech", {}) .get("speech", {})
@@ -50,43 +50,46 @@ class LLMProvider(LLMProviderBase):
# 用于处理跨chunk的标签 # 用于处理跨chunk的标签
buffer = "" buffer = ""
for chunk in responses: try:
try: for chunk in responses:
delta = ( try:
chunk.choices[0].delta delta = (
if getattr(chunk, "choices", None) chunk.choices[0].delta
else None if getattr(chunk, "choices", None)
) else None
content = delta.content if hasattr(delta, "content") else "" )
content = delta.content if hasattr(delta, "content") else ""
if content: if content:
# 将内容添加到缓冲区 # 将内容添加到缓冲区
buffer += content buffer += content
# 处理缓冲区中的标签 # 处理缓冲区中的标签
while "<think>" in buffer and "</think>" in buffer: while "<think>" in buffer and "</think>" in buffer:
# 找到完整的<think></think>标签并移除 # 找到完整的<think></think>标签并移除
pre = buffer.split("<think>", 1)[0] pre = buffer.split("<think>", 1)[0]
post = buffer.split("</think>", 1)[1] post = buffer.split("</think>", 1)[1]
buffer = pre + post buffer = pre + post
# 处理只有开始标签的情况 # 处理只有开始标签的情况
if "<think>" in buffer: if "<think>" in buffer:
is_active = False is_active = False
buffer = buffer.split("<think>", 1)[0] buffer = buffer.split("<think>", 1)[0]
# 处理只有结束标签的情况 # 处理只有结束标签的情况
if "</think>" in buffer: if "</think>" in buffer:
is_active = True is_active = True
buffer = buffer.split("</think>", 1)[1] buffer = buffer.split("</think>", 1)[1]
# 如果当前处于活动状态且缓冲区有内容,则输出 # 如果当前处于活动状态且缓冲区有内容,则输出
if is_active and buffer: if is_active and buffer:
yield buffer yield buffer
buffer = "" # 清空缓冲区 buffer = "" # 清空缓冲区
except Exception as e: except Exception as e:
logger.bind(tag=TAG).error(f"Error processing chunk: {e}") logger.bind(tag=TAG).error(f"Error processing chunk: {e}")
finally:
responses.close()
def response_with_functions(self, session_id, dialogue, functions=None): def response_with_functions(self, session_id, dialogue, functions=None):
# 如果是qwen3模型,在用户最后一条消息中添加/no_think指令 # 如果是qwen3模型,在用户最后一条消息中添加/no_think指令
@@ -117,49 +120,52 @@ class LLMProvider(LLMProviderBase):
is_active = True is_active = True
buffer = "" buffer = ""
for chunk in stream: try:
try: for chunk in stream:
delta = ( try:
chunk.choices[0].delta delta = (
if getattr(chunk, "choices", None) chunk.choices[0].delta
else None if getattr(chunk, "choices", None)
) else None
content = delta.content if hasattr(delta, "content") else None )
tool_calls = ( content = delta.content if hasattr(delta, "content") else None
delta.tool_calls if hasattr(delta, "tool_calls") else None tool_calls = (
) delta.tool_calls if hasattr(delta, "tool_calls") else None
)
# 如果是工具调用,直接传递 # 如果是工具调用,直接传递
if tool_calls: if tool_calls:
yield None, tool_calls yield None, tool_calls
continue
# 处理文本内容
if content:
# 将内容添加到缓冲区
buffer += content
# 处理缓冲区中的标签
while "<think>" in buffer and "</think>" in buffer:
# 找到完整的<think></think>标签并移除
pre = buffer.split("<think>", 1)[0]
post = buffer.split("</think>", 1)[1]
buffer = pre + post
# 处理只有开始标签的情况
if "<think>" in buffer:
is_active = False
buffer = buffer.split("<think>", 1)[0]
# 处理只有结束标签的情况
if "</think>" in buffer:
is_active = True
buffer = buffer.split("</think>", 1)[1]
# 如果当前处于活动状态且缓冲区有内容,则输出
if is_active and buffer:
yield buffer, None
buffer = "" # 清空缓冲区
except Exception as e:
logger.bind(tag=TAG).error(f"Error processing function chunk: {e}")
continue continue
finally:
# 处理文本内容 stream.close()
if content:
# 将内容添加到缓冲区
buffer += content
# 处理缓冲区中的标签
while "<think>" in buffer and "</think>" in buffer:
# 找到完整的<think></think>标签并移除
pre = buffer.split("<think>", 1)[0]
post = buffer.split("</think>", 1)[1]
buffer = pre + post
# 处理只有开始标签的情况
if "<think>" in buffer:
is_active = False
buffer = buffer.split("<think>", 1)[0]
# 处理只有结束标签的情况
if "</think>" in buffer:
is_active = True
buffer = buffer.split("</think>", 1)[1]
# 如果当前处于活动状态且缓冲区有内容,则输出
if is_active and buffer:
yield buffer, None
buffer = "" # 清空缓冲区
except Exception as e:
logger.bind(tag=TAG).error(f"Error processing function chunk: {e}")
continue
@@ -115,21 +115,24 @@ class LLMProvider(LLMProviderBase):
responses = self.client.chat.completions.create(**request_params) responses = self.client.chat.completions.create(**request_params)
is_active = True is_active = True
for chunk in responses: try:
try: for chunk in responses:
delta = chunk.choices[0].delta if getattr(chunk, "choices", None) else None try:
content = getattr(delta, "content", "") if delta else "" delta = chunk.choices[0].delta if getattr(chunk, "choices", None) else None
except IndexError: content = getattr(delta, "content", "") if delta else ""
content = "" except IndexError:
if content: content = ""
if "<think>" in content: if content:
is_active = False if "<think>" in content:
content = content.split("<think>")[0] is_active = False
if "</think>" in content: content = content.split("<think>")[0]
is_active = True if "</think>" in content:
content = content.split("</think>")[-1] is_active = True
if is_active: content = content.split("</think>")[-1]
yield content if is_active:
yield content
finally:
responses.close()
def response_with_functions(self, session_id, dialogue, functions=None, **kwargs): def response_with_functions(self, session_id, dialogue, functions=None, **kwargs):
dialogue = self.normalize_dialogue(dialogue) dialogue = self.normalize_dialogue(dialogue)
@@ -157,16 +160,19 @@ class LLMProvider(LLMProviderBase):
stream = self.client.chat.completions.create(**request_params) stream = self.client.chat.completions.create(**request_params)
for chunk in stream: try:
if getattr(chunk, "choices", None): for chunk in stream:
delta = chunk.choices[0].delta if getattr(chunk, "choices", None):
content = getattr(delta, "content", "") delta = chunk.choices[0].delta
tool_calls = getattr(delta, "tool_calls", None) content = getattr(delta, "content", "")
yield content, tool_calls tool_calls = getattr(delta, "tool_calls", None)
elif isinstance(getattr(chunk, "usage", None), CompletionUsage): yield content, tool_calls
usage_info = getattr(chunk, "usage", None) elif isinstance(getattr(chunk, "usage", None), CompletionUsage):
logger.bind(tag=TAG).info( usage_info = getattr(chunk, "usage", None)
f"Token 消耗:输入 {getattr(usage_info, 'prompt_tokens', '未知')}" logger.bind(tag=TAG).info(
f"输出 {getattr(usage_info, 'completion_tokens', '未知')}" f"Token 消耗:输入 {getattr(usage_info, 'prompt_tokens', '未知')}"
f"共计 {getattr(usage_info, 'total_tokens', '未知')}" f"输出 {getattr(usage_info, 'completion_tokens', '未知')}"
) f"共计 {getattr(usage_info, 'total_tokens', '未知')}"
)
finally:
stream.close()
@@ -74,12 +74,15 @@ class LLMProvider(LLMProviderBase):
tools=functions, tools=functions,
) )
for chunk in stream: try:
delta = chunk.choices[0].delta for chunk in stream:
content = delta.content delta = chunk.choices[0].delta
tool_calls = delta.tool_calls content = delta.content
tool_calls = delta.tool_calls
if content: if content:
yield content, tool_calls yield content, tool_calls
elif tool_calls: elif tool_calls:
yield None, tool_calls yield None, tool_calls
finally:
stream.close()