diff --git a/src/agent/middlewares/retry_on_rate_limit.py b/src/agent/middlewares/retry_on_rate_limit.py index 1414570..fc7d190 100644 --- a/src/agent/middlewares/retry_on_rate_limit.py +++ b/src/agent/middlewares/retry_on_rate_limit.py @@ -2,15 +2,34 @@ from __future__ import annotations import asyncio +import queue +import time from typing import Any from langchain.agents.middleware import AgentMiddleware, AgentState -from langchain_core.messages import ToolMessage _PAUSE = 30 # секунд ожидания при 429 _TRIES = 5 # максимум попыток +# Глобальный канал событий для UI (устанавливается из ui.py перед запуском агента). +# Если None — события просто не отправляются (CLI-режим). +_ui_event_queue: queue.Queue | None = None + + +def set_ui_event_queue(q: queue.Queue | None) -> None: + """Вызывается из ui.py чтобы подключить очередь событий.""" + global _ui_event_queue + _ui_event_queue = q + + +def _emit(event: dict) -> None: + if _ui_event_queue is not None: + try: + _ui_event_queue.put_nowait(event) + except Exception: + pass + def _is_429(exc: Exception) -> bool: msg = str(exc) @@ -18,7 +37,9 @@ def _is_429(exc: Exception) -> bool: class RetryOnRateLimitMiddleware(AgentMiddleware[AgentState[Any], Any]): - """Перехватывает 429 от любого инструмента и повторяет с паузой.""" + """Перехватывает 429 от любого инструмента и повторяет с паузой. + Отправляет события rate_limit_wait / rate_limit_retry в UI-очередь. + """ def wrap_tool_call(self, request, handler): for attempt in range(1, _TRIES + 1): @@ -28,8 +49,12 @@ class RetryOnRateLimitMiddleware(AgentMiddleware[AgentState[Any], Any]): if _is_429(e) and attempt < _TRIES: name = request.tool_call.get("name", "") print(f"[retry-mw] {name} → 429, жду {_PAUSE}с (попытка {attempt}/{_TRIES})...") - import time + _emit({"t": "rate_limit_wait", "name": name, + "pause": _PAUSE, "attempt": attempt, "max": _TRIES, + "ts": _now()}) time.sleep(_PAUSE) + _emit({"t": "rate_limit_retry", "name": name, + "attempt": attempt + 1, "ts": _now()}) else: raise @@ -41,6 +66,16 @@ class RetryOnRateLimitMiddleware(AgentMiddleware[AgentState[Any], Any]): if _is_429(e) and attempt < _TRIES: name = request.tool_call.get("name", "") print(f"[retry-mw] {name} → 429, жду {_PAUSE}с (попытка {attempt}/{_TRIES})...") + _emit({"t": "rate_limit_wait", "name": name, + "pause": _PAUSE, "attempt": attempt, "max": _TRIES, + "ts": _now()}) await asyncio.sleep(_PAUSE) + _emit({"t": "rate_limit_retry", "name": name, + "attempt": attempt + 1, "ts": _now()}) else: raise + + +def _now() -> str: + from datetime import datetime + return datetime.now().strftime("%H:%M:%S") diff --git a/ui.py b/ui.py index bf31c71..bbcb50b 100644 --- a/ui.py +++ b/ui.py @@ -147,14 +147,22 @@ class AgentCallback(BaseCallbackHandler): # Запуск агента в фоне # --------------------------------------------------------------------------- -def _run_agent_thread(agent, messages, config, q: queue.Queue, cb: AgentCallback): +def _run_agent_thread( + agent, messages, config, q: queue.Queue, cb: AgentCallback, + stop_flag: threading.Event, +): + from src.agent.middlewares.retry_on_rate_limit import set_ui_event_queue + set_ui_event_queue(q) async def _inner(): try: result = await agent.ainvoke(messages, {**config, "callbacks": [cb]}) q.put({"t": "done", "result": result}) except Exception as e: q.put({"t": "fatal", "msg": str(e)}) - asyncio.run(_inner()) + finally: + set_ui_event_queue(None) + if not stop_flag.is_set(): + asyncio.run(_inner()) # --------------------------------------------------------------------------- @@ -207,6 +215,21 @@ def _render_event(ev: dict) -> str: msg = ev["msg"].replace("<", "<") return f'