fix(llm): make stream cancellation responsive
This commit is contained in:
parent
b830dbfeec
commit
0a01a94c93
|
|
@ -8,6 +8,8 @@
|
|||
|
||||
## Unreleased
|
||||
|
||||
- 修复国际旗舰模型输出思考过程时偶尔无法及时停止的问题；即使模型暂时没有返回新的流式片段，点击“停止”也会迅速中断当前回答。
|
||||
|
||||
- Windows Node 现在可在“专业软件”区域自动检测或手工指定应用位置，并分别安装、更新 Origin、ANSYS、Blender 的独立运行环境；只装有部分专业软件的节点不再需要处理无关软件。
|
||||
|
||||
- 对话中尚未发送的文字现在会按对话自动暂存；助手回答或后台进程执行期间也可以继续发送补充信息，消息会显示为待处理并在当前工作结束后依次发送。动作按钮采用单按钮形态：忙碌且输入为空时用于停止，输入文字或加入附件后自动切回发送。切换对话或刷新页面后，草稿和待处理消息仍会保留。
|
||||
|
|
|
|||
115
core/llm.py
115
core/llm.py
|
|
@ -2,14 +2,16 @@
|
|||
|
||||
`chat()`:同步阻塞,一次性返回完整 response。给 probe / 离线探测用。
|
||||
`chat_stream()`:流式 generator,yield chunk;调用方累积 + 用 litellm.stream_chunk_builder
|
||||
拼回完整 response。loop 走这条以便 chunk 之间 poll cancel(同步 LLM call 不可中断;
|
||||
流式下 cancel 延迟 ~ chunk 间隔 100ms 级,而非整轮 generation 时长几十秒)。
|
||||
拼回完整 response。loop 传 cancel_check 后由独立 pump 承担 provider 阻塞读取，
|
||||
控制线程固定节拍响应停止，不依赖下一块到达。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import queue
|
||||
import threading
|
||||
import time
|
||||
from typing import Any, Iterator, List, Optional
|
||||
from typing import Any, Callable, Iterator, List, Optional
|
||||
|
||||
# 跳过启动时从 GitHub 拉 model_prices 的网络请求,直接用 litellm 打包的本地副本。
|
||||
# 必须在 `import litellm` 之前设置,否则 get_model_cost_map() 已经跑过了。
|
||||
|
|
@ -104,17 +106,34 @@ class LLM:
|
|||
parallel_tool_calls: Optional[bool] = None,
|
||||
reasoning_effort: Optional[str] = None,
|
||||
max_retries: int = 3,
|
||||
cancel_check: Optional[Callable[[], bool]] = None,
|
||||
) -> Iterator[Any]:
|
||||
"""流式 chat:yield 每个 chunk。调用方累积 + 用 litellm.stream_chunk_builder 拼回完整 response。
|
||||
|
||||
重试语义:连接建立阶段错误(还没拿到第一个 chunk)按 max_retries 退避重试;
|
||||
开始流之后失败直接抛(半截 partial 没法续)。usage 通过 stream_options.include_usage
|
||||
让最后一个 chunk 带 usage。
|
||||
|
||||
传入 cancel_check 时，底层阻塞读取在 daemon pump 线程执行，当前 run 线程按
|
||||
100ms 轮询取消。这样 provider 在 TTFT / reasoning 分片之间长时间无字节时，
|
||||
停止也不必等到下一块到达；已有底层 stream 会被主动 close。
|
||||
"""
|
||||
kwargs = self._build_kwargs(messages, tools, parallel_tool_calls, reasoning_effort)
|
||||
kwargs["stream"] = True
|
||||
kwargs["stream_options"] = {"include_usage": True}
|
||||
|
||||
if cancel_check is not None:
|
||||
yield from self._chat_stream_interruptible(
|
||||
kwargs, max_retries=max_retries, cancel_check=cancel_check,
|
||||
)
|
||||
return
|
||||
|
||||
yield from self._chat_stream_direct(kwargs, max_retries=max_retries)
|
||||
|
||||
@staticmethod
|
||||
def _chat_stream_direct(kwargs: dict, *, max_retries: int) -> Iterator[Any]:
|
||||
"""无取消调用方的直连路径（probe 等保持原有同步语义）。"""
|
||||
|
||||
last_err: Optional[Exception] = None
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
|
|
@ -139,3 +158,93 @@ class LLM:
|
|||
close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def _chat_stream_interruptible(
|
||||
kwargs: dict,
|
||||
*,
|
||||
max_retries: int,
|
||||
cancel_check: Callable[[], bool],
|
||||
) -> Iterator[Any]:
|
||||
"""把 provider 的阻塞迭代与 run 控制线程解耦，以固定节拍响应取消。"""
|
||||
items: queue.Queue[tuple[str, Any]] = queue.Queue(maxsize=64)
|
||||
stop = threading.Event()
|
||||
stream_box: dict[str, Any] = {}
|
||||
|
||||
def offer(kind: str, value: Any = None) -> bool:
|
||||
while not stop.is_set():
|
||||
try:
|
||||
items.put((kind, value), timeout=0.1)
|
||||
return True
|
||||
except queue.Full:
|
||||
continue
|
||||
return False
|
||||
|
||||
def pump() -> None:
|
||||
stream = None
|
||||
try:
|
||||
last_err: Optional[Exception] = None
|
||||
for attempt in range(max_retries):
|
||||
if stop.is_set():
|
||||
return
|
||||
try:
|
||||
stream = litellm.completion(**kwargs)
|
||||
stream_box["stream"] = stream
|
||||
break
|
||||
except (
|
||||
RateLimitError,
|
||||
APIConnectionError,
|
||||
ServiceUnavailableError,
|
||||
Timeout,
|
||||
APIError,
|
||||
) as e:
|
||||
last_err = e
|
||||
if attempt == max_retries - 1:
|
||||
raise
|
||||
if stop.wait(2 ** attempt):
|
||||
return
|
||||
else:
|
||||
if last_err is not None:
|
||||
raise last_err
|
||||
return
|
||||
|
||||
for chunk in stream:
|
||||
if not offer("chunk", chunk):
|
||||
return
|
||||
offer("done")
|
||||
except BaseException as e: # noqa: BLE001 - 原样转抛到 run 线程
|
||||
offer("error", e)
|
||||
finally:
|
||||
if stream is not None:
|
||||
close = getattr(stream, "close", None)
|
||||
if callable(close):
|
||||
try:
|
||||
close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
worker = threading.Thread(target=pump, daemon=True, name="llm-stream-pump")
|
||||
worker.start()
|
||||
try:
|
||||
while True:
|
||||
if cancel_check():
|
||||
return
|
||||
try:
|
||||
kind, value = items.get(timeout=0.1)
|
||||
except queue.Empty:
|
||||
continue
|
||||
if kind == "chunk":
|
||||
yield value
|
||||
elif kind == "error":
|
||||
raise value
|
||||
else:
|
||||
return
|
||||
finally:
|
||||
stop.set()
|
||||
stream = stream_box.get("stream")
|
||||
close = getattr(stream, "close", None)
|
||||
if callable(close):
|
||||
try:
|
||||
close()
|
||||
except Exception:
|
||||
pass
|
||||
|
|
|
|||
|
|
@ -614,6 +614,7 @@ class AgentLoop:
|
|||
messages=llm_messages,
|
||||
tools=self.executor.schemas(),
|
||||
reasoning_effort=self.caps.default_reasoning_effort or None,
|
||||
cancel_check=self._is_cancelled,
|
||||
)
|
||||
cancelled = False
|
||||
try:
|
||||
|
|
@ -633,6 +634,10 @@ class AgentLoop:
|
|||
delta_reasoning = extract_delta_reasoning(chunk)
|
||||
if delta_reasoning:
|
||||
self._emit({"type": "reasoning", "delta": delta_reasoning})
|
||||
# interruptible stream 会在无新 chunk 的等待期直接因 cancel 结束迭代；
|
||||
# 循环体没有机会执行上面的检查，故在正常耗尽处再判一次。
|
||||
if self._is_cancelled():
|
||||
cancelled = True
|
||||
finally:
|
||||
# generator 提前 break 时 GeneratorExit 触发 chat_stream finally → close 底层连接
|
||||
close = getattr(stream, "close", None)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,63 @@
|
|||
"""流式 LLM 在 provider 无新分片时仍能快速响应用户停止。"""
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
import time
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
from core.llm import LLM
|
||||
|
||||
|
||||
class _BlockingStream:
|
||||
def __init__(self) -> None:
|
||||
self.reading = threading.Event()
|
||||
self.released = threading.Event()
|
||||
self.closed = False
|
||||
|
||||
def __iter__(self):
|
||||
return self
|
||||
|
||||
def __next__(self):
|
||||
self.reading.set()
|
||||
self.released.wait(30)
|
||||
raise StopIteration
|
||||
|
||||
def close(self) -> None:
|
||||
self.closed = True
|
||||
self.released.set()
|
||||
|
||||
|
||||
class LlmStreamCancelTests(unittest.TestCase):
|
||||
def test_interruptible_stream_preserves_normal_chunks(self) -> None:
|
||||
llm = object.__new__(LLM)
|
||||
llm._build_kwargs = lambda *args, **kwargs: {"model": "test"}
|
||||
with patch("core.llm.litellm.completion", return_value=iter(["a", "b"])):
|
||||
chunks = list(llm.chat_stream([], cancel_check=lambda: False))
|
||||
self.assertEqual(chunks, ["a", "b"])
|
||||
|
||||
def test_cancel_does_not_wait_for_next_provider_chunk(self) -> None:
|
||||
llm = object.__new__(LLM)
|
||||
llm._build_kwargs = lambda *args, **kwargs: {"model": "test"}
|
||||
raw = _BlockingStream()
|
||||
cancelled = threading.Event()
|
||||
|
||||
def trigger_cancel() -> None:
|
||||
self.assertTrue(raw.reading.wait(2))
|
||||
cancelled.set()
|
||||
|
||||
trigger = threading.Thread(target=trigger_cancel, daemon=True)
|
||||
trigger.start()
|
||||
started = time.monotonic()
|
||||
with patch("core.llm.litellm.completion", return_value=raw):
|
||||
chunks = list(llm.chat_stream([], cancel_check=cancelled.is_set))
|
||||
elapsed = time.monotonic() - started
|
||||
trigger.join(2)
|
||||
|
||||
self.assertEqual(chunks, [])
|
||||
self.assertTrue(raw.closed)
|
||||
self.assertLess(elapsed, 1.0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Loading…
Reference in New Issue