Merge pull request #214 from YooooEX/telegram-polish
fix(tgapp): handle Telegram stream rate limits
This commit is contained in:
1 file changed
+183
-19
+183
-19
@@ -5,6 +5,7 @@ from agentmain import GeneraticAgent
|
||||
try:
|
||||
from telegram import BotCommand, InlineKeyboardButton, InlineKeyboardMarkup
|
||||
from telegram.constants import ChatType, MessageLimit, ParseMode
|
||||
from telegram.error import RetryAfter
|
||||
from telegram.ext import ApplicationBuilder, CallbackQueryHandler, MessageHandler, filters, ContextTypes
|
||||
from telegram.helpers import escape_markdown
|
||||
from telegram.request import HTTPXRequest
|
||||
@@ -34,6 +35,9 @@ ALLOWED = set(mykeys.get('tg_allowed_users', []))
|
||||
_DRAFT_HINT = "thinking..."
|
||||
_STREAM_SUFFIX = " ⏳"
|
||||
_STREAM_SEGMENT_LIMIT = max(1200, MessageLimit.MAX_TEXT_LENGTH - 256)
|
||||
_STREAM_UPDATE_INTERVAL_SECONDS = 2.0
|
||||
_STREAM_MIN_UPDATE_CHARS = 400
|
||||
_RETRY_AFTER_MARGIN_SECONDS = 1.0
|
||||
_QUEUE_WAIT_SECONDS = 1
|
||||
_ASK_USER_HOOK_KEY = "telegram_ask_user_menu"
|
||||
_ASK_CALLBACK_PREFIX = "ask:"
|
||||
@@ -69,7 +73,41 @@ def _make_draft_id():
|
||||
|
||||
def _visible_segments(text):
|
||||
text = (text or "").strip()
|
||||
return split_text(text, _STREAM_SEGMENT_LIMIT) if text else []
|
||||
if not text:
|
||||
return []
|
||||
segments = []
|
||||
for part in split_text(text, _STREAM_SEGMENT_LIMIT):
|
||||
segments.extend(_markdown_safe_segments(part))
|
||||
return segments
|
||||
|
||||
def _markdown_safe_segments(text, limit=None):
|
||||
limit = limit or MessageLimit.MAX_TEXT_LENGTH
|
||||
text = (text or "").strip()
|
||||
if not text:
|
||||
return []
|
||||
if len(_to_markdown_v2(text)) <= limit:
|
||||
return [text]
|
||||
parts = []
|
||||
remaining = text
|
||||
while remaining:
|
||||
if len(_to_markdown_v2(remaining)) <= limit:
|
||||
parts.append(remaining)
|
||||
break
|
||||
low, high, best = 1, len(remaining), 1
|
||||
while low <= high:
|
||||
mid = (low + high) // 2
|
||||
if len(_to_markdown_v2(remaining[:mid].rstrip() or remaining[:mid])) <= limit:
|
||||
best = mid
|
||||
low = mid + 1
|
||||
else:
|
||||
high = mid - 1
|
||||
cut = remaining.rfind("\n", 0, best)
|
||||
if cut < max(1, best * 0.6):
|
||||
cut = best
|
||||
chunk = remaining[:cut].rstrip() or remaining[:best]
|
||||
parts.append(chunk)
|
||||
remaining = remaining[len(chunk):].lstrip()
|
||||
return parts
|
||||
|
||||
def _line_complete(line):
|
||||
return (line or "").endswith(("\n", "\r"))
|
||||
@@ -341,12 +379,77 @@ class _TelegramStreamSession:
|
||||
self.files = []
|
||||
self.sent_segments = 0
|
||||
self.active_display = ""
|
||||
self.pending_display = ""
|
||||
self.retry_until = 0.0
|
||||
self.last_update_at = 0.0
|
||||
self.last_update_raw_len = 0
|
||||
|
||||
def _now(self):
|
||||
return time.monotonic()
|
||||
|
||||
def _retry_after_seconds(self, exc):
|
||||
retry_after = getattr(exc, "_retry_after", None)
|
||||
if retry_after is None:
|
||||
retry_after = getattr(exc, "retry_after", 0) or 0
|
||||
if hasattr(retry_after, "total_seconds"):
|
||||
retry_after = retry_after.total_seconds()
|
||||
try:
|
||||
return max(0.0, float(retry_after))
|
||||
except (TypeError, ValueError):
|
||||
return 0.0
|
||||
|
||||
def _set_retry_after(self, exc):
|
||||
wait_seconds = self._retry_after_seconds(exc) + _RETRY_AFTER_MARGIN_SECONDS
|
||||
self.retry_until = max(self.retry_until, self._now() + wait_seconds)
|
||||
|
||||
def _is_retrying(self):
|
||||
return self._now() < self.retry_until
|
||||
|
||||
async def _wait_for_retry(self):
|
||||
remaining = self.retry_until - self._now()
|
||||
if remaining > 0:
|
||||
await asyncio.sleep(remaining)
|
||||
|
||||
def _should_stream_update(self, display):
|
||||
if display == self.active_display:
|
||||
return False
|
||||
if self.last_update_at <= 0:
|
||||
return True
|
||||
elapsed = self._now() - self.last_update_at
|
||||
raw_delta = len(self.raw_text) - self.last_update_raw_len
|
||||
return elapsed >= _STREAM_UPDATE_INTERVAL_SECONDS or raw_delta >= _STREAM_MIN_UPDATE_CHARS
|
||||
|
||||
def _mark_stream_update(self, display):
|
||||
self.active_display = display
|
||||
self.pending_display = ""
|
||||
self.last_update_at = self._now()
|
||||
self.last_update_raw_len = len(self.raw_text)
|
||||
|
||||
def _stream_display(self, text):
|
||||
base = (text or _DRAFT_HINT).strip() or _DRAFT_HINT
|
||||
safe_parts = _markdown_safe_segments(base)
|
||||
base = safe_parts[-1] if safe_parts else _DRAFT_HINT
|
||||
if base == _DRAFT_HINT:
|
||||
return base
|
||||
display = base + _STREAM_SUFFIX
|
||||
if len(_to_markdown_v2(display)) <= MessageLimit.MAX_TEXT_LENGTH:
|
||||
return display
|
||||
return base
|
||||
|
||||
async def prime(self):
|
||||
if self.can_use_draft and await self._send_draft(_DRAFT_HINT):
|
||||
if self.can_use_draft:
|
||||
draft_result = await self._send_draft(_DRAFT_HINT)
|
||||
if draft_result is True:
|
||||
self.active_display = _DRAFT_HINT
|
||||
return
|
||||
if draft_result is None:
|
||||
self.active_display = _DRAFT_HINT
|
||||
return
|
||||
try:
|
||||
await self._upsert_live_message(_DRAFT_HINT, wait_retry=False)
|
||||
except RetryAfter:
|
||||
self.active_display = _DRAFT_HINT
|
||||
return
|
||||
await self._upsert_live_message(_DRAFT_HINT)
|
||||
self.active_display = _DRAFT_HINT
|
||||
|
||||
async def add_chunk(self, chunk):
|
||||
@@ -395,16 +498,24 @@ class _TelegramStreamSession:
|
||||
await self._stream_active(active_text)
|
||||
|
||||
async def _stream_active(self, text):
|
||||
display = (text or _DRAFT_HINT).strip() or _DRAFT_HINT
|
||||
if display != _DRAFT_HINT:
|
||||
display = display + _STREAM_SUFFIX
|
||||
display = self._stream_display(text)
|
||||
if display == self.active_display:
|
||||
return
|
||||
if self.can_use_draft and await self._send_draft(display):
|
||||
self.active_display = display
|
||||
self.pending_display = display
|
||||
if self._is_retrying() or not self._should_stream_update(display):
|
||||
return
|
||||
try:
|
||||
if self.can_use_draft:
|
||||
draft_result = await self._send_draft(display)
|
||||
if draft_result is True:
|
||||
self._mark_stream_update(display)
|
||||
return
|
||||
if draft_result is None:
|
||||
return
|
||||
await self._upsert_live_message(display, wait_retry=False)
|
||||
self._mark_stream_update(display)
|
||||
except RetryAfter:
|
||||
return
|
||||
await self._upsert_live_message(display)
|
||||
self.active_display = display
|
||||
|
||||
async def _finalize_segment(self, text):
|
||||
final_text = (text or "").strip() or "..."
|
||||
@@ -428,6 +539,9 @@ class _TelegramStreamSession:
|
||||
parse_mode=ParseMode.MARKDOWN_V2,
|
||||
)
|
||||
return True
|
||||
except RetryAfter as exc:
|
||||
self._set_retry_after(exc)
|
||||
return None
|
||||
except Exception as exc:
|
||||
if _is_not_modified_error(exc):
|
||||
return True
|
||||
@@ -436,30 +550,71 @@ class _TelegramStreamSession:
|
||||
self.draft_id = _make_draft_id()
|
||||
return False
|
||||
|
||||
async def _reply_text(self, text):
|
||||
async def _retry_call(self, func, *args):
|
||||
while True:
|
||||
await self._wait_for_retry()
|
||||
try:
|
||||
return await func(*args)
|
||||
except RetryAfter as exc:
|
||||
self._set_retry_after(exc)
|
||||
|
||||
async def _reply_text_once(self, text):
|
||||
markdown = _to_markdown_v2(text)
|
||||
try:
|
||||
return await self.root_msg.reply_text(markdown, parse_mode=ParseMode.MARKDOWN_V2)
|
||||
except RetryAfter as exc:
|
||||
self._set_retry_after(exc)
|
||||
raise
|
||||
except Exception as exc:
|
||||
if _is_not_modified_error(exc):
|
||||
return None
|
||||
return await self.root_msg.reply_text(text)
|
||||
try:
|
||||
return await self.root_msg.reply_text(text)
|
||||
except RetryAfter as retry_exc:
|
||||
self._set_retry_after(retry_exc)
|
||||
raise
|
||||
|
||||
async def _edit_text(self, msg, text):
|
||||
async def _reply_text(self, text, wait_retry=True):
|
||||
last_msg = None
|
||||
for segment in _markdown_safe_segments(text) or ["..."]:
|
||||
if wait_retry:
|
||||
last_msg = await self._retry_call(self._reply_text_once, segment)
|
||||
else:
|
||||
last_msg = await self._reply_text_once(segment)
|
||||
return last_msg
|
||||
|
||||
async def _edit_text_once(self, msg, text):
|
||||
markdown = _to_markdown_v2(text)
|
||||
try:
|
||||
updated = await msg.edit_text(markdown, parse_mode=ParseMode.MARKDOWN_V2)
|
||||
except RetryAfter as exc:
|
||||
self._set_retry_after(exc)
|
||||
raise
|
||||
except Exception as exc:
|
||||
if _is_not_modified_error(exc):
|
||||
return msg
|
||||
updated = await msg.edit_text(text)
|
||||
try:
|
||||
updated = await msg.edit_text(text)
|
||||
except RetryAfter as retry_exc:
|
||||
self._set_retry_after(retry_exc)
|
||||
raise
|
||||
return updated if hasattr(updated, "edit_text") else msg
|
||||
|
||||
async def _upsert_live_message(self, text):
|
||||
if self.live_msg is None:
|
||||
self.live_msg = await self._reply_text(text)
|
||||
async def _edit_text(self, msg, text, wait_retry=True):
|
||||
segments = _markdown_safe_segments(text) or ["..."]
|
||||
if wait_retry:
|
||||
updated = await self._retry_call(self._edit_text_once, msg, segments[0])
|
||||
else:
|
||||
self.live_msg = await self._edit_text(self.live_msg, text)
|
||||
updated = await self._edit_text_once(msg, segments[0])
|
||||
for segment in segments[1:]:
|
||||
updated = await self._reply_text(segment, wait_retry=wait_retry)
|
||||
return updated if hasattr(updated, "edit_text") else msg
|
||||
|
||||
async def _upsert_live_message(self, text, wait_retry=True):
|
||||
if self.live_msg is None:
|
||||
self.live_msg = await self._reply_text(text, wait_retry=wait_retry)
|
||||
else:
|
||||
self.live_msg = await self._edit_text(self.live_msg, text, wait_retry=wait_retry)
|
||||
|
||||
|
||||
class _TelegramTurnStreamCoordinator:
|
||||
@@ -578,9 +733,18 @@ async def _stream(dq, msg):
|
||||
break
|
||||
except asyncio.CancelledError:
|
||||
await stream.finish_with_notice("⏹️ 已停止")
|
||||
except RetryAfter as exc:
|
||||
print(f"[TG stream retry_after] {type(exc).__name__}: {exc}", flush=True)
|
||||
if stream.session is not None:
|
||||
stream.session._set_retry_after(exc)
|
||||
except Exception as exc:
|
||||
print(f"[TG stream error] {type(exc).__name__}: {exc}", flush=True)
|
||||
await stream.finish_with_notice(f"❌ 输出失败: {exc}")
|
||||
if stream.session is not None and stream.session._is_retrying():
|
||||
return
|
||||
try:
|
||||
await stream.finish_with_notice(f"❌ 输出失败: {exc}")
|
||||
except RetryAfter as retry_exc:
|
||||
print(f"[TG stream error notice retry_after] {type(retry_exc).__name__}: {retry_exc}", flush=True)
|
||||
|
||||
def _normalized_command(text):
|
||||
parts = (text or "").strip().split(None, 1)
|
||||
|
||||
Reference in new issue
Block a user