diff --git a/CHANGELOG.md b/CHANGELOG.md index f690aa4..fa7227f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,8 @@ ## Unreleased +- 修复国际旗舰模型输出思考过程时偶尔无法及时停止的问题；即使模型暂时没有返回新的流式片段，点击“停止”也会迅速中断当前回答。 + - Windows Node 现在可在“专业软件”区域自动检测或手工指定应用位置，并分别安装、更新 Origin、ANSYS、Blender 的独立运行环境；只装有部分专业软件的节点不再需要处理无关软件。 - 对话中尚未发送的文字现在会按对话自动暂存；助手回答或后台进程执行期间也可以继续发送补充信息，消息会显示为待处理并在当前工作结束后依次发送。动作按钮采用单按钮形态：忙碌且输入为空时用于停止，输入文字或加入附件后自动切回发送。切换对话或刷新页面后，草稿和待处理消息仍会保留。 diff --git a/core/llm.py b/core/llm.py index fa213a3..9703ec8 100644 --- a/core/llm.py +++ b/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 diff --git a/core/loop.py b/core/loop.py index 71673c1..f2e2cc2 100644 --- a/core/loop.py +++ b/core/loop.py @@ -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) diff --git a/tests/test_llm_stream_cancel.py b/tests/test_llm_stream_cancel.py new file mode 100644 index 0000000..a589915 --- /dev/null +++ b/tests/test_llm_stream_cancel.py @@ -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()