diff --git a/README.md b/README.md index d8cceb2..9b1d57a 100644 --- a/README.md +++ b/README.md @@ -155,6 +155,8 @@ Each is off by default. The holding-news DM and the movers reply are switched on the last 3 days are not repeated. Each message has an **Unsubscribe** button: per ticker, or all holding news, with an **Undo** button afterwards. People must have messaged the bot (or a group it is in) once so it knows their user ID. `/holdingnews` in a DM runs the check now and always replies. + If the AI provider is at capacity, the daily check is retried every 30 minutes for up to 12 hours + (nothing is sent to anyone meanwhile). - **Reply to a movers bot** (switched on by listing the bots in `MOVERS_BOTS`; empty, the default, means off): when such a bot posts an end-of-day list ("≥ 5.0% at close (ASX):" followed by stocks with % changes) the bot explains each move with a forced search, once, without tagging the other bot, ignoring its image, diff --git a/lib/config.py b/lib/config.py index fdf6509..a2e7ff1 100644 --- a/lib/config.py +++ b/lib/config.py @@ -55,6 +55,8 @@ SEARCH_TOKENS = 400 # cap on each search write-up (OpenRouter search runner) MAX_SEARCHES = 5 # search calls run per round MAX_TG_MESSAGE = 4000 # Telegram's limit is 4096 +HOLDING_NEWS_RETRY_EVERY = 30 * 60 # seconds between retries of a holding-news check the provider was too busy for +HOLDING_NEWS_RETRY_FOR = 12 * 3600 # seconds after which it gives up class ConfigError(Exception): diff --git a/lib/features/holding_news.py b/lib/features/holding_news.py index 93a42c2..fe561e0 100644 --- a/lib/features/holding_news.py +++ b/lib/features/holding_news.py @@ -15,8 +15,10 @@ from telegram.error import Forbidden from telegram.ext import ContextTypes +from .. import config from ..ask import ask from ..context import Ctx +from ..llm.base import is_busy from ..llm.gate import Route from ..mcp.results import table_records from ..prompts import HOLDING_NEWS_PROMPT @@ -113,26 +115,28 @@ async def send_holding_news(ctx: Ctx, chat_id: int, news: str, codes: list[str]) await asyncio.to_thread(ctx.store.remember_news_message, chat_id, sent.message_id, codes) -async def run_holding_news(ctx: Ctx, only_username: str | None = None - ) -> dict[str, tuple[str, list[str]] | None]: +async def run_holding_news(ctx: Ctx, only_username: str | None = None, *, usernames: set[str] | None = None, + busy: list[str] | None = None) -> dict[str, tuple[str, list[str]] | None]: """Check each recipient's holdings; DM anyone with major news. - Returns {username: (news text, codes it covers) or None}.""" + `usernames` limits the scheduled run to those people; `busy` collects the ones the provider was + too busy to check. Returns {username: (news text, codes it covers) or None}.""" server = ctx.registry.servers.get("sharesight") if not (server and server.session): log.warning("Holding news: Sharesight isn't connected, skipping") return {} people: dict[str, list[str]] = {} for portfolio, username in ctx.st.holding_news_recipients.items(): - if only_username is None or username == only_username: + if ((only_username is None or username == only_username) + and (usernames is None or username in usernames)): people.setdefault(username, []).append(portfolio) results: dict[str, tuple[str, list[str]] | None] = {} for username, portfolio_names in people.items(): - results.update(await _check_person(ctx, server, username, portfolio_names, only_username)) + results.update(await _check_person(ctx, server, username, portfolio_names, only_username, busy)) return results async def _check_person(ctx: Ctx, server, username: str, portfolio_names: list[str], - only_username: str | None) -> dict: + only_username: str | None, busy: list[str] | None = None) -> dict: muted = await asyncio.to_thread(ctx.store.muted_codes, username) if "*" in muted: log.info("Holding news: @%s has unsubscribed from all of it", username) @@ -143,8 +147,14 @@ async def _check_person(ctx: Ctx, server, username: str, portfolio_names: list[s if not holdings: return {} news = await holding_news_text(ctx, username, holdings) - except Exception: - log.exception("Holding news check failed for @%s", username) + except Exception as e: + if is_busy(e): # automated, so no chat message; one log line is enough + log.warning("Holding news check failed for @%s, the provider is busy: %s", + username, " ".join(str(e).split())) + if busy is not None: + busy.append(username) + else: + log.exception("Holding news check failed for @%s", username) return {} result = (news, news_codes(news, holdings)) if news else None log.info("Holding news for @%s: %s", username, "found" if news else "nothing major") @@ -179,13 +189,32 @@ async def daily_loop(ctx: Ctx) -> None: wait = seconds_until(ctx.st.holding_news_time, datetime.now(ctx.st.tz)) log.info("Next holding news check in %.1f hours", wait / 3600) await asyncio.sleep(wait) + busy: list[str] = [] try: - await run_holding_news(ctx) + await run_holding_news(ctx, busy=busy) + await retry_busy(ctx, busy) except Exception: log.exception("Holding news check failed") await asyncio.sleep(1) +async def retry_busy(ctx: Ctx, usernames: list[str]) -> None: + """Check again, for just the people the provider was too busy for, every HOLDING_NEWS_RETRY_EVERY + seconds until it works or HOLDING_NEWS_RETRY_FOR has passed.""" + waited = 0 + while usernames and waited < config.HOLDING_NEWS_RETRY_FOR: + log.info("Holding news: trying again for %s in %d min", ", ".join(f"@{u}" for u in usernames), + config.HOLDING_NEWS_RETRY_EVERY // 60) + await asyncio.sleep(config.HOLDING_NEWS_RETRY_EVERY) + waited += config.HOLDING_NEWS_RETRY_EVERY + busy: list[str] = [] + await run_holding_news(ctx, usernames=set(usernames), busy=busy) + usernames = busy + if usernames: + log.warning("Holding news: gave up on %s after %d hours", ", ".join(f"@{u}" for u in usernames), + config.HOLDING_NEWS_RETRY_FOR // 3600) + + def username_of(msg_or_user) -> str: return (getattr(msg_or_user, "username", None) or "").lower() diff --git a/lib/llm/base.py b/lib/llm/base.py index 0ad329d..15f72da 100644 --- a/lib/llm/base.py +++ b/lib/llm/base.py @@ -8,6 +8,17 @@ from ..mcp.schema import ToolDef +BUSY_STATUSES = (429, 503, 529) + + +class Overloaded(RuntimeError): + """The provider is out of capacity right now (an error inside a failed response).""" + + +def is_busy(e: BaseException) -> bool: + """Capacity or rate limiting: the person can simply try again in a few minutes.""" + return isinstance(e, Overloaded) or getattr(e, "status_code", None) in BUSY_STATUSES + @dataclass class Usage: diff --git a/lib/llm/policy.py b/lib/llm/policy.py index 19caf5a..c6968db 100644 --- a/lib/llm/policy.py +++ b/lib/llm/policy.py @@ -8,7 +8,7 @@ from openai import APIStatusError from ..mcp.server import Registry -from .base import Answer, Backend, Request +from .base import Answer, Backend, Request, is_busy from .runner import run log = logging.getLogger("bot") @@ -22,7 +22,8 @@ async def ask(backend: Backend, registry: Registry, req: Request, *, dedupe: boo try: return await run(backend, registry, req, dedupe=dedupe) except (APIStatusError, RuntimeError) as first: - retriable = isinstance(first, RuntimeError) or first.status_code in (400, 422) + retriable = (isinstance(first, RuntimeError) and not is_busy(first)) or ( + isinstance(first, APIStatusError) and first.status_code in (400, 422)) if req.must_search or not retriable or not (req.search or req.tools): raise log.warning("Retrying without tools: %s", first) diff --git a/lib/llm/xai.py b/lib/llm/xai.py index 27e05c2..7e5c6b9 100644 --- a/lib/llm/xai.py +++ b/lib/llm/xai.py @@ -7,7 +7,7 @@ from openai import AsyncOpenAI, BadRequestError from ..config import Settings -from .base import Backend, Call, Request, Step, Usage +from .base import Backend, Call, Overloaded, Request, Step, Usage log = logging.getLogger("bot") @@ -33,6 +33,18 @@ def _as_input(item) -> dict: return item.model_dump(exclude_none=True) if hasattr(item, "model_dump") else dict(item) +def answer_text(resp) -> str: + """The response's last message. `output_text` joins every message, and with server-side search + the model can write a draft, search, then write the answer, all in one response.""" + messages = [i for i in resp.output if i.type == "message"] + texts = ["".join(getattr(c, "text", "") or "" for c in m.content if c.type == "output_text") + for m in messages] + if len([t for t in texts if t.strip()]) > 1: + log.info("xAI response had %d messages; kept the last (items: %s)", len(texts), + " ".join(i.type for i in resp.output)) + return next((t for t in reversed(texts) if t.strip()), resp.output_text or "").strip() + + class XaiBackend(Backend): name = "xai" search_what = "the web and X (Twitter)" @@ -74,7 +86,9 @@ async def _create(self, req: Request, kw: dict, **inputs): stream = await self.client.responses.create(**kw, **inputs) text, final = "", None async for event in stream: - if event.type == "response.output_text.delta": + if event.type == "response.output_item.added" and event.item.type == "message": + text = "" # a new message replaces the draft before it + elif event.type == "response.output_text.delta": text += event.delta if req.on_text: req.on_text(text) @@ -83,7 +97,8 @@ async def _create(self, req: Request, kw: dict, **inputs): if final is None: raise RuntimeError("The stream ended without a final response") if final.status == "failed": - raise RuntimeError(f"xAI response failed: {final.error}") + msg = f"xAI response failed: {final.error}" + raise (Overloaded if getattr(final.error, "code", None) == "server_error" else RuntimeError)(msg) return final async def step(self, conv: dict, req: Request, *, tool_choice: str | None) -> Step: @@ -112,7 +127,7 @@ async def step(self, conv: dict, req: Request, *, tool_choice: str | None) -> St cost=(getattr(u, "cost_in_usd_ticks", 0) or 0) / TICKS_PER_USD, ) searches = sum(1 for i in resp.output if i.type.endswith("_search_call")) - return Step((resp.output_text or "").strip(), calls, usage, resp.status, searches) + return Step(answer_text(resp), calls, usage, resp.status, searches) def add_results(self, conv: dict, step: Step, results: list[str]) -> None: items = [{"type": "function_call_output", "call_id": c.id, "output": out} diff --git a/lib/telegram/handlers.py b/lib/telegram/handlers.py index fd6ba8a..ad131b6 100644 --- a/lib/telegram/handlers.py +++ b/lib/telegram/handlers.py @@ -17,6 +17,7 @@ from ..context import Ctx from ..features import dm_buttons, holding_news, movers, periods, post from ..history import compact, format_rows +from ..llm.base import is_busy from ..llm.gate import Route, is_portfolio_question, route from ..msgtext import describe, sender_name from ..prompts import MOVERS_EXPLAIN, PRESET_PROMPT, chat_prompt, down_note @@ -29,7 +30,8 @@ log = logging.getLogger("bot") PORTFOLIO_ADMINS_ONLY = "Portfolio data is for admins only." -FOLLOW_UP_WORDS = 12 # a message this short right after the bot's reply may be a follow-up to it +BUSY_REPLY = "The AI provider is at capacity right now. Please try again in a few minutes." +FOLLOW_UP_WORDS = 12 # a message this short right after the bot's reply may be a follow-up to it FOLLOW_UP_SECONDS = 30 * 60 FOLLOW_UP_CHARS = 400 # of the bot's reply that the tool gate also reads READ_IMAGE_WORDS = re.compile(r"\b(read|text|says?|chart|table|number|screenshot|ocr|zoom|detail|label)\b", re.I) @@ -206,9 +208,13 @@ async def _answer(self, msg: Message, trig: Trigger) -> None: try: html_text = await self._generate(msg, trig, draft) except Exception as e: - log.exception("Reply failed") - html_text = ("Sorry, something went wrong:\n" - f"
{html.escape(f'{type(e).__name__}: {e}'[:1500])}
") + if is_busy(e): + log.warning("Reply failed, the provider is busy: %s", " ".join(str(e).split())) + html_text = BUSY_REPLY + else: + log.exception("Reply failed") + html_text = ("Sorry, something went wrong:\n" + f"
{html.escape(f'{type(e).__name__}: {e}'[:1500])}
") finally: if draft: draft.stop() diff --git a/tests/test_features.py b/tests/test_features.py index afc607a..d7b5120 100644 --- a/tests/test_features.py +++ b/tests/test_features.py @@ -9,6 +9,7 @@ from lib import config from lib.context import Ctx from lib.features import dm_buttons, holding_news, movers +from lib.llm.base import Overloaded from lib.telegram import commands from .conftest import FakeBot, make_msg, update_for, user_msg @@ -79,6 +80,32 @@ async def test_digest_says_nothing_when_nothing_and_respects_mutes(env, store): assert "AMD (NASDAQ)" in prompt and "NVDA (NASDAQ)" not in prompt # muted code left out +async def test_busy_provider_is_retried_for_just_that_person_then_delivered(env, store, monkeypatch, caplog): + monkeypatch.setattr(config, "HOLDING_NEWS_RETRY_EVERY", 0) + ctx = holding_ctx(env, store, [Overloaded("at capacity"), Overloaded("at capacity"), step("• NVDA: big news")]) + busy: list[str] = [] + assert await holding_news.run_holding_news(ctx, busy=busy) == {} and busy == ["bob_llama"] + assert ctx.bot.sent == [] + await holding_news.retry_busy(ctx, busy) # still busy once, then works + assert len(ctx.backend.seen) == 3 and "Holding news" in ctx.bot.sent[-1]["text"] + assert not any(r.exc_info for r in caplog.records) + + +async def test_busy_provider_is_given_up_on_after_the_retry_window(env, store, monkeypatch, caplog): + monkeypatch.setattr(config, "HOLDING_NEWS_RETRY_EVERY", 4 * 3600) + monkeypatch.setattr(config, "HOLDING_NEWS_RETRY_FOR", 12 * 3600) + slept = [] + + async def no_sleep(seconds): + slept.append(seconds) + + monkeypatch.setattr(holding_news.asyncio, "sleep", no_sleep) + ctx = holding_ctx(env, store, [Overloaded("at capacity")] * 4) + await holding_news.retry_busy(ctx, ["bob_llama"]) + assert slept == [4 * 3600] * 3 and len(ctx.backend.seen) == 3 and ctx.bot.sent == [] # 12 h / 4 h + assert "gave up on @bob_llama after 12 hours" in caplog.text + + async def test_digest_skips_people_the_bot_cannot_reach(env, store): ctx = holding_ctx(env, store, [step("• NVDA: big news")]) ctx.bot.fail_send[5] = Forbidden("bot was blocked") diff --git a/tests/test_handlers.py b/tests/test_handlers.py index 1b2c487..ba7f5a5 100644 --- a/tests/test_handlers.py +++ b/tests/test_handlers.py @@ -5,7 +5,7 @@ from lib import config from lib.context import Ctx from lib.features import dm_buttons -from lib.llm.base import Usage +from lib.llm.base import Overloaded, Usage from lib.telegram.handlers import Handlers, format_answer, mentions from lib.telegram.handlers import photo_file_id as pick_photo @@ -178,6 +178,16 @@ async def test_error_becomes_a_visible_message(env, store): assert "Sorry, something went wrong" in m.replies[0]["text"] and "provider exploded" in m.replies[0]["text"] +async def test_busy_provider_gets_a_plain_reply_and_a_one_line_log(env, store, caplog): + err = Overloaded("xAI response failed: ResponseError(code='server_error',\n message='at capacity')") + ctx, backend, h = make_ctx(env, store, [err]) + m = user_msg(ctx.bot, "@stockbot hello") + await run(h, m) + assert "at capacity right now" in m.replies[0]["text"] and "
" not in m.replies[0]["text"]
+    line = next(r.getMessage() for r in caplog.records if "provider is busy" in r.getMessage())
+    assert "\n" not in line and not any(r.exc_info for r in caplog.records)
+
+
 async def test_movers_list_off_by_default_then_answered_once_without_tagging(env, store):
     text = "≥ 5.0% at close (ASX):\nNVDA +6.1%\nAMD +5.5%"
     for flag, expect in (({}, 0), ({"MOVERS_BOTS": "finbotibot"}, 1)):
diff --git a/tests/test_llm_backends.py b/tests/test_llm_backends.py
index 91a7623..79a25b2 100644
--- a/tests/test_llm_backends.py
+++ b/tests/test_llm_backends.py
@@ -5,7 +5,7 @@
 from openai import BadRequestError
 
 from lib import config
-from lib.llm.base import Usage
+from lib.llm.base import Overloaded, Usage
 from lib.llm.openrouter import OpenRouterBackend
 from lib.llm.xai import XaiBackend
 from lib.llm.zai import ZaiBackend
@@ -225,6 +225,41 @@ async def test_xai_cache_headers_and_flat_tools(env):
     assert kw["input"][0] == {"role": "system", "content": "SYS"} and kw["reasoning"] == {"effort": "low"}
 
 
+def xai_message(text):
+    from types import SimpleNamespace as N
+    return N(type="message", content=[N(type="output_text", text=text)])
+
+
+async def test_xai_keeps_only_the_last_message_when_it_drafts_before_searching():
+    from types import SimpleNamespace as N
+    st = config.load({"TELEGRAM_BOT_TOKEN": "1:x", "XAI_API_KEY": "k"})
+    items = [xai_message("draft"), N(type="web_search_call"), xai_message("answer")]
+    events = [N(type="response.output_item.added", item=items[0]),
+              N(type="response.output_text.delta", delta="draft"),
+              N(type="response.output_item.added", item=items[2]),
+              N(type="response.output_text.delta", delta="answer"),
+              xai_final(items, text="draftanswer")]
+    seen = []
+    b = XaiBackend(st, FakeXaiClient(events))
+    r = req(on_text=seen.append)
+    step = await b.step(b.start(r), r, tool_choice=None)
+    assert step.text == "answer" and step.searches == 1
+    assert seen == ["draft", "answer"]
+
+
+async def test_xai_failed_response_at_capacity_is_overloaded_other_failures_are_not():
+    from types import SimpleNamespace as N
+    st = config.load({"TELEGRAM_BOT_TOKEN": "1:x", "XAI_API_KEY": "k"})
+    for code, exc in (("server_error", Overloaded), ("invalid_prompt", RuntimeError)):
+        final = xai_final(status="failed")
+        final.response.error = N(code=code, message="m")
+        b = XaiBackend(st, FakeXaiClient([final]))
+        r = req()
+        with pytest.raises(exc) as e:
+            await b.step(b.start(r), r, tool_choice=None)
+        assert isinstance(e.value, Overloaded) == (exc is Overloaded)
+
+
 async def test_xai_cost_comes_from_usd_ticks(env):
     from types import SimpleNamespace as N
     st = config.load({"TELEGRAM_BOT_TOKEN": "1:x", "XAI_API_KEY": "k"})
diff --git a/tests/test_runner.py b/tests/test_runner.py
index 9e775b5..17b3d38 100644
--- a/tests/test_runner.py
+++ b/tests/test_runner.py
@@ -3,7 +3,7 @@
 from openai import APIStatusError, BadRequestError
 
 from lib import config
-from lib.llm.base import Usage
+from lib.llm.base import Overloaded, Usage
 from lib.llm.gate import route, wants_tools
 from lib.llm.policy import ask
 from lib.llm.runner import call_query, run, run_calls
@@ -198,6 +198,10 @@ async def test_no_retry_for_auth_credit_rate_limit_server_errors_or_must_search(
     with pytest.raises(RuntimeError, match="x"):
         await ask(b, registry(), req(must_search=True, search=True))
     assert len(b.seen) == 1
+    b = ScriptedBackend([Overloaded("at capacity"), step("never")])  # tools aren't the cause
+    with pytest.raises(Overloaded):
+        await ask(b, registry(mcp), with_tools(mcp))
+    assert len(b.seen) == 1
     b = ScriptedBackend([RuntimeError("x"), step("never")])  # nothing to take away: no retry
     with pytest.raises(RuntimeError):
         await ask(b, registry(), req())