File size: 2,802 Bytes
ad88834 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | import asyncio
import logging
import math
import time
from collections.abc import Awaitable, Callable
from contextlib import suppress
from agentcade.expressions import segmented_progress_vector
from observability import describe
logger = logging.getLogger(__name__)
class ProgressDisplay:
def __init__(self, show: Callable[[str], Awaitable[None]], lock: asyncio.Lock):
self.show = show
self.lock = lock
self.task: asyncio.Task | None = None
self.token = 0
self.active = False
self.step = 0
async def frame(self, token: int, step: int) -> None:
async with self.lock:
if self.active and token == self.token:
await self.show(segmented_progress_vector(step))
self.step = step
async def start(self) -> int:
await self.cancel()
self.active = True
self.step = 0
await self.frame(self.token, 0)
return self.token
def advance(self, token: int, seconds: float) -> None:
if not math.isfinite(seconds) or seconds <= 0:
raise ValueError("预计生成时间必须是有限正数")
if not self.active or token != self.token:
return
if self.task is not None:
raise RuntimeError("当前进度任务已经开始计时")
self.task = asyncio.create_task(self.animate(token, seconds))
async def animate(self, token: int, seconds: float) -> None:
started = time.monotonic()
while self.active and token == self.token and self.step < 9:
deadline = started + (self.step + 1) * seconds / 10
await asyncio.sleep(max(0, deadline - time.monotonic()))
step = min(9, int((time.monotonic() - started) * 10 / seconds))
await self.frame(token, step)
async def stop_timer(self) -> None:
task, self.task = self.task, None
if task is not None:
task.cancel()
try:
with suppress(asyncio.CancelledError):
await task
except Exception as error:
# 进度画面只是提示，显示失败不影响后面的输出。
logger.warning("进度显示失败：%s", describe(error), exc_info=True)
async def finish(self) -> None:
await self.stop_timer()
if self.active:
await self.frame(self.token, 10)
self.active = False
self.token += 1
await asyncio.sleep(0.25)
async def cancel(self) -> None:
self.active = False
self.token += 1
await self.stop_timer()
async with self.lock:
pass
async def check(self) -> None:
if self.task is not None and self.task.done():
await self.stop_timer()
|