diff --git a/src/agent/profiles/react_to_me.py b/src/agent/profiles/react_to_me.py index c119736..e9037c5 100644 --- a/src/agent/profiles/react_to_me.py +++ b/src/agent/profiles/react_to_me.py @@ -36,7 +36,10 @@ def __init__( embedding: Embeddings, ) -> None: super().__init__(llm, embedding) - self.llm = llm + # A streaming copy, the same way the RAG chains get theirs. The UI + # displays only what the callback handler streams, so an answer + # produced without streaming is invisible however correct it is. + self.live_llm = llm.model_copy(update={"streaming": True}) self.unsafe_answer_generator: Runnable = create_unsafe_answer_generator( llm, streaming=True @@ -149,7 +152,9 @@ async def preprocess( active_sources=active_sources, ) - async def _answer_from_live_services(self, state: ReactToMeState) -> ReactToMeState: + async def _answer_from_live_services( + self, state: ReactToMeState, config: RunnableConfig + ) -> ReactToMeState: """Answer from the MCP tools instead of the vector store. If the server is unreachable the question falls back to retrieval, with @@ -166,17 +171,19 @@ async def _answer_from_live_services(self, state: ReactToMeState) -> ReactToMeSt ) fallback = dict(state) fallback["active_sources"] = ["reactome"] - return await self.generate_answer( - ReactToMeState(**fallback), RunnableConfig() - ) + # The real config, not a fresh one: it carries the callbacks the UI + # streams through, and a fallback nobody can see is not a fallback. + return await self.generate_answer(ReactToMeState(**fallback), config) answer = await answer_from_live_services( # BaseChatModel satisfies ToolCallingModel at runtime; its # bind_tools signature is wider than the Protocol restates. - cast(ToolCallingModel, self.llm), + cast(ToolCallingModel, self.live_llm), tools, state["rephrased_input"], language=state["detected_language"], + chat_history=state["chat_history"] or None, + config=config, ) return ReactToMeState( chat_history=[HumanMessage(state["user_input"]), AIMessage(answer)], @@ -207,7 +214,7 @@ async def generate_answer( ) -> ReactToMeState: source = state["active_sources"][0] if source == "live": - return await self._answer_from_live_services(state) + return await self._answer_from_live_services(state, config) rag = self.rags[source] result: dict[str, Any] = await rag.ainvoke( { diff --git a/src/reactome_mcp/answer.py b/src/reactome_mcp/answer.py index b2a181f..daf95b1 100644 --- a/src/reactome_mcp/answer.py +++ b/src/reactome_mcp/answer.py @@ -12,6 +12,7 @@ """ import logging +from collections.abc import Sequence from typing import Any, Protocol, runtime_checkable from langchain_core.messages import ( @@ -20,6 +21,7 @@ SystemMessage, ToolMessage, ) +from langchain_core.runnables import RunnableConfig from langchain_core.tools import BaseTool logger = logging.getLogger(__name__) @@ -60,9 +62,16 @@ async def answer_from_live_services( tools: list[BaseTool], question: str, language: str = "English", - chat_history: list[BaseMessage] | None = None, + chat_history: Sequence[BaseMessage] | None = None, + config: RunnableConfig | None = None, ) -> str: - """Run the tool-calling loop and return the answer text.""" + """Run the tool-calling loop and return the answer text. + + `config` is not optional in practice. It carries the callback handlers, and + the Chainlit UI displays only what those handlers stream -- `chat-chainlit.py` + reads `chainlit_cb.final_stream` and has no path that posts the returned + string. Without it this function answers correctly into the void. + """ by_name = {tool.name: tool for tool in tools} bound = llm.bind_tools(tools) @@ -73,7 +82,7 @@ async def answer_from_live_services( ] for _round in range(MAX_TOOL_ROUNDS): - reply = await bound.ainvoke(messages) + reply = await bound.ainvoke(messages, config) messages.append(reply) calls = getattr(reply, "tool_calls", None) or [] @@ -111,7 +120,7 @@ async def answer_from_live_services( "Answer now, from the tool results above. Do not call more tools." ) ) - reply = await llm.ainvoke(messages) + reply = await llm.ainvoke(messages, config) content: Any = reply.content if isinstance(content, list): diff --git a/tests/reactome_mcp/test_live_answer.py b/tests/reactome_mcp/test_live_answer.py index ffc5f75..1713adb 100644 --- a/tests/reactome_mcp/test_live_answer.py +++ b/tests/reactome_mcp/test_live_answer.py @@ -36,6 +36,7 @@ def __init__(self, replies: list[AIMessage]) -> None: self._replies = list(replies) self.seen: list[list[Any]] = [] self.bound_tools: list[Any] = [] + self.configs: list[Any] = [] def bind_tools(self, tools: Any, **kwargs: Any) -> "_FakeLLM": self.bound_tools = tools @@ -43,6 +44,7 @@ def bind_tools(self, tools: Any, **kwargs: Any) -> "_FakeLLM": async def ainvoke(self, messages: Any, *args: Any, **kwargs: Any) -> AIMessage: self.seen.append(list(messages)) + self.configs.append(args[0] if args else kwargs.get("config")) return self._replies.pop(0) if self._replies else AIMessage("out of replies") @@ -152,3 +154,67 @@ def test_content_blocks_are_flattened() -> None: llm = _FakeLLM([AIMessage([{"type": "text", "text": "96 species."}])]) answer = asyncio.run(answer_from_live_services(llm, [reactome_species], "x")) assert answer == "96 species." + + +def test_the_config_reaches_every_model_call() -> None: + """Without this the answer is correct and invisible. + + The Chainlit UI displays only what the callback handlers stream -- + `chat-chainlit.py` reads `chainlit_cb.final_stream` and has no path that + posts the returned string. The callbacks travel in the RunnableConfig, so a + config that is dropped anywhere in this loop means the user sees nothing at + all while the tests pass. + + Shipped exactly that way on 2026-09-15: measured 0 streamed tokens for a + live answer, against 103 after the fix. + """ + config = {"callbacks": ["sentinel"]} + llm = _FakeLLM( + [ + AIMessage("", tool_calls=[_call("reactome_species")]), + AIMessage("Reactome covers 96 species."), + ] + ) + asyncio.run( + answer_from_live_services(llm, [reactome_species], "x", config=config) # type: ignore[arg-type] + ) + + assert llm.configs, "the model was never invoked" + assert all( + c == config for c in llm.configs + ), f"a model call lost the config: {llm.configs}" + + +def test_the_config_reaches_the_forced_final_answer_too() -> None: + """The round cap takes a different path out of the loop, and it needs the + config just as much -- that call is the one that produces the text.""" + config = {"callbacks": ["sentinel"]} + llm = _FakeLLM( + [ + AIMessage("", tool_calls=[_call("reactome_species", str(i))]) + for i in range(10) + ] + ) + asyncio.run( + answer_from_live_services(llm, [reactome_species], "x", config=config) # type: ignore[arg-type] + ) + + assert all(c == config for c in llm.configs) + + +def test_chat_history_is_carried_into_the_conversation() -> None: + """A follow-up such as "can you tell me?" is meaningless without it.""" + from langchain_core.messages import HumanMessage + + history = [HumanMessage("what species are in reactome")] + llm = _FakeLLM([AIMessage("96 species.")]) + asyncio.run( + answer_from_live_services( + llm, [reactome_species], "can you tell me?", chat_history=history + ) + ) + + sent = llm.seen[0] + assert any( + isinstance(m, HumanMessage) and "what species" in str(m.content) for m in sent + ), "the history never reached the model"