118 lines
3.0 KiB
Python
118 lines
3.0 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
import logging
|
||
|
|
from collections.abc import Callable
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
_DEFAULT_LABELS = [
|
||
|
|
"思考中...",
|
||
|
|
"处理中...",
|
||
|
|
"分析中...",
|
||
|
|
"检索中...",
|
||
|
|
"生成中...",
|
||
|
|
]
|
||
|
|
_DEFAULT_INITIAL_DELAY_MS = 5000
|
||
|
|
|
||
|
|
|
||
|
|
class ProgressDraftGate:
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
on_start: Callable[[], None],
|
||
|
|
initial_delay_ms: float = _DEFAULT_INITIAL_DELAY_MS,
|
||
|
|
):
|
||
|
|
self._on_start = on_start
|
||
|
|
self._initial_delay = initial_delay_ms / 1000.0
|
||
|
|
self._started = False
|
||
|
|
self._disposed = False
|
||
|
|
self._work_events = 0
|
||
|
|
self._timer: asyncio.Task | None = None
|
||
|
|
|
||
|
|
async def note_work(self) -> bool:
|
||
|
|
if self._disposed:
|
||
|
|
return False
|
||
|
|
self._work_events += 1
|
||
|
|
if self._started:
|
||
|
|
return True
|
||
|
|
if self._work_events > 1:
|
||
|
|
await self._start()
|
||
|
|
return True
|
||
|
|
self._schedule()
|
||
|
|
return False
|
||
|
|
|
||
|
|
async def _start(self) -> None:
|
||
|
|
if self._started or self._disposed:
|
||
|
|
return
|
||
|
|
self._started = True
|
||
|
|
self._cancel_timer()
|
||
|
|
self._on_start()
|
||
|
|
|
||
|
|
def _schedule(self) -> None:
|
||
|
|
if self._timer or self._started or self._disposed:
|
||
|
|
return
|
||
|
|
self._timer = asyncio.create_task(self._delayed_start())
|
||
|
|
|
||
|
|
async def _delayed_start(self) -> None:
|
||
|
|
await asyncio.sleep(self._initial_delay)
|
||
|
|
await self._start()
|
||
|
|
|
||
|
|
def cancel(self) -> None:
|
||
|
|
self._disposed = True
|
||
|
|
self._cancel_timer()
|
||
|
|
|
||
|
|
def _cancel_timer(self) -> None:
|
||
|
|
if self._timer and not self._timer.done():
|
||
|
|
self._timer.cancel()
|
||
|
|
self._timer = None
|
||
|
|
|
||
|
|
|
||
|
|
def format_tool_progress_line(
|
||
|
|
tool_name: str,
|
||
|
|
status: str = "",
|
||
|
|
labels: list[str] | None = None,
|
||
|
|
) -> str:
|
||
|
|
emoji_map = {
|
||
|
|
"search": "🔍",
|
||
|
|
"read": "📖",
|
||
|
|
"write": "✏️",
|
||
|
|
"exec": "🛠️",
|
||
|
|
"bash": "🛠️",
|
||
|
|
"shell": "🛠️",
|
||
|
|
"web_search": "🌐",
|
||
|
|
"fetch": "📡",
|
||
|
|
"think": "💭",
|
||
|
|
"analyze": "📊",
|
||
|
|
}
|
||
|
|
emoji = emoji_map.get(tool_name, "🔧")
|
||
|
|
label_map = {
|
||
|
|
"search": "搜索",
|
||
|
|
"read": "读取",
|
||
|
|
"write": "写入",
|
||
|
|
"exec": "执行",
|
||
|
|
"bash": "执行命令",
|
||
|
|
"web_search": "网络搜索",
|
||
|
|
"fetch": "获取数据",
|
||
|
|
"think": "思考",
|
||
|
|
"analyze": "分析",
|
||
|
|
}
|
||
|
|
label = label_map.get(tool_name, tool_name)
|
||
|
|
|
||
|
|
if status:
|
||
|
|
return f"{emoji} {label}: {status}"
|
||
|
|
return f"{emoji} {label}"
|
||
|
|
|
||
|
|
|
||
|
|
class ProgressDraftRenderer:
|
||
|
|
def __init__(self, labels: list[str] | None = None):
|
||
|
|
self._labels = labels or _DEFAULT_LABELS
|
||
|
|
self._index = 0
|
||
|
|
|
||
|
|
def next_label(self) -> str:
|
||
|
|
label = self._labels[self._index % len(self._labels)]
|
||
|
|
self._index += 1
|
||
|
|
return label
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def format_tool_line(tool_name: str, status: str = "") -> str:
|
||
|
|
return format_tool_progress_line(tool_name, status)
|