diff --git a/lib/llm/gate.py b/lib/llm/gate.py index b80dbc0..6f05d10 100644 --- a/lib/llm/gate.py +++ b/lib/llm/gate.py @@ -34,6 +34,9 @@ LIVE_WORDS = re.compile( r"\b(news|today|latest|yesterday|tonight|this week|right now|breaking|announce\w*|who won|" r"score|weather|happening|headlines?|current\w*|updates?|recent\w*|just (now|in))\b", re.I) +# Asking about a trading session ("what happened at close"): answers must be fetched, not recalled. +SESSION_WORDS = re.compile( + r"\b(?:(?:at|the|into|after) (?:the )?(?:close|open)|pre-?market|after[- ]hours|closing|opening)\b", re.I) @dataclass(frozen=True) @@ -65,6 +68,10 @@ def is_portfolio_question(text: str, st: Settings) -> bool: return bool(PORTFOLIO_WEAK.search(text)) and not _has_ticker(text) +def is_session_question(text: str) -> bool: + return bool(SESSION_WORDS.search(text)) + + def route(text: str, st: Settings, registry: Registry, *, force_full: bool = False) -> Route: """Tools for this message. With TOKEN_SAVER off, or `force_full` (a search/data request such as a movers list, a preset button or a retry after NEEDS_TOOLS), everything is offered.""" @@ -73,7 +80,7 @@ def route(text: str, st: Settings, registry: Registry, *, force_full: bool = Fal return Route(up, st.search) portfolio = is_portfolio_question(text, st) market = portfolio or bool(MARKET_WORDS.search(text)) or _has_ticker(text) - live = bool(LIVE_WORDS.search(text)) + live = is_session_question(text) or bool(LIVE_WORDS.search(text)) servers = [s for s in up if (portfolio if s.cfg.get("gate") == "portfolio" else market)] search = st.search and (market or live) simple = not (servers or search) diff --git a/lib/prompts.py b/lib/prompts.py index d9cb3da..dcfda7a 100644 --- a/lib/prompts.py +++ b/lib/prompts.py @@ -69,7 +69,7 @@ def down_note(down: list[MCPServer]) -> str: def chat_prompt(*, transcript: str, sender: str, private: bool, reply_quote: str | None, msg_id: int, now: str, down: str = "", saver: bool = True, middle: str = "", - fresh: bool = False) -> str: + fresh: bool = False, session: bool = False) -> str: """The user turn: transcript, the replied-to message, and the volatile tail (kept last so the cached prefix before it is unchanged).""" head = (f"Private one-to-one chat with {sender}, not the group. Transcript:\n{transcript}" @@ -89,6 +89,9 @@ def chat_prompt(*, transcript: str, sender: str, private: bool, reply_quote: str "for anything else asked, such as prices or a comparison with market data); don't reuse " "numbers or refusals from earlier replies. When comparing portfolios, or asked \"why\", " "write 2-3 short paragraphs on the tailwinds and headwinds of each portfolio, not a summary.") + if session: # a trading-session question: news and reasons in earlier replies are stale by now + tail += (" Search for current news before answering, and say why each stock moved. Don't reuse " + "news or reasons from earlier messages in the transcript.") return head + tail + down diff --git a/lib/telegram/handlers.py b/lib/telegram/handlers.py index b4f669e..8ab1550 100644 --- a/lib/telegram/handlers.py +++ b/lib/telegram/handlers.py @@ -18,7 +18,7 @@ from ..features import dm_buttons, holding_news, movers, periods, post from ..history import compact, format_rows from ..llm.base import BUSY_REPLY, is_busy -from ..llm.gate import Route, is_portfolio_question, route +from ..llm.gate import Route, is_portfolio_question, is_session_question, route from ..msgtext import describe, sender_name from ..prompts import MOVERS_EXPLAIN, PRESET_PROMPT, chat_prompt, down_note from ..store import is_dm @@ -179,7 +179,8 @@ async def _gate_text(self, msg: Message, trig: Trigger) -> str: text += " " + compact(prev.text)[:FOLLOW_UP_CHARS] return text - async def _chat_text(self, msg: Message, trig: Trigger, middle: str = "", fresh: bool = False) -> str: + async def _chat_text(self, msg: Message, trig: Trigger, middle: str = "", fresh: bool = False, + session: bool = False) -> str: ctx, st = self.ctx, self.ctx.st rows = await asyncio.to_thread(ctx.store.history, msg.chat_id, st.history_limit) transcript = format_rows(rows, ctx.self_name, st.tz, line_max=st.history_line_max, @@ -187,7 +188,8 @@ async def _chat_text(self, msg: Message, trig: Trigger, middle: str = "", fresh: return chat_prompt( transcript=transcript, sender=sender_name(msg), private=trig.private, reply_quote=self._quote(trig), msg_id=msg.message_id, now=ctx.now(), - down=down_note(ctx.registry.down()), saver=st.token_saver, middle=middle, fresh=fresh) + down=down_note(ctx.registry.down()), saver=st.token_saver, middle=middle, fresh=fresh, + session=session) # -- answering --------------------------------------------------------------------- async def _answer(self, msg: Message, trig: Trigger) -> None: @@ -257,9 +259,11 @@ async def _generate(self, msg: Message, trig: Trigger, draft: Draft | None) -> s if chosen.simple: # the model must not be told it has no tools chosen = Route([], False) fresh = is_portfolio_question(context, st) and any(s.cfg.get("gate") == "portfolio" for s in chosen.servers) - parts = [{"type": "text", "text": await self._chat_text(msg, trig, fresh=fresh)}] + session = is_session_question(context) and not chosen.simple + parts = [{"type": "text", "text": await self._chat_text(msg, trig, fresh=fresh, session=session)}] parts += await self._images(msg, trig) answer = await ask(ctx, parts, chosen, + must_search=session and chosen.search, # a session recap must not be recalled cache_key=f"chat-{chat_id}", on_text=on_text, chat_id=chat_id, extra_tools=extra, admin=admin) return format_answer(answer.text, movers_bots=st.movers_bots, summary=trig.movers) diff --git a/tests/test_handlers.py b/tests/test_handlers.py index ba7f5a5..e31610c 100644 --- a/tests/test_handlers.py +++ b/tests/test_handlers.py @@ -90,6 +90,14 @@ async def test_partial_route_that_needs_more_is_asked_again_with_everything(env, assert "NEEDS_TOOLS" in backend.seen[0]["system"] # the partial route told the model how to ask +async def test_session_question_forces_a_search(env, store): + ctx, backend, h = make_ctx(env, store, [step("fresh recap", searches=1)], servers=[FakeMcp("yahoo")]) + m = user_msg(ctx.bot, "@stockbot what happened at close?") + await run(h, m) + assert backend.seen[0]["choice"] == "required" and backend.seen[0]["search"] is True + assert "Search for current news" in backend.seen[0]["prompt"] + + async def test_market_question_offers_tools_up_front(env, store): mcp = FakeMcp("yahoo") ctx, backend, h = make_ctx(env, store, [step("ok")], servers=[mcp]) diff --git a/tests/test_runner.py b/tests/test_runner.py index 17b3d38..5c6e16d 100644 --- a/tests/test_runner.py +++ b/tests/test_runner.py @@ -242,6 +242,22 @@ def test_gate_live_question_gets_search_only(settings): assert r.search and not r.servers +@pytest.mark.parametrize("text", ["what happened at close?", "how did it go at the open", "any pre-market moves"]) +def test_gate_market_session_words_are_live(text, settings): + y, s = servers() + r = route(text, settings, registry(y, s)) + assert r.search and not r.simple + + +def test_session_question_prompt_demands_current_news(): + from lib.llm.gate import is_session_question + from lib.prompts import chat_prompt + assert is_session_question("what happened at close?") and not is_session_question("hello") + kw = dict(transcript="t", sender="A", private=False, reply_quote=None, msg_id=1, now="now") + assert "Search for current news" in chat_prompt(**kw, session=True) + assert "Search for current news" not in chat_prompt(**kw) + + def test_gate_full_when_forced_or_token_saver_off(settings, env): y, s = servers() assert route("hello", settings, registry(y, s), force_full=True).servers == [y, s]