From bdb2b3c21be5fdd15f828b5831bd2df55db6dfc2 Mon Sep 17 00:00:00 2001 From: kiranfin <4finjas@online.de> Date: Wed, 7 Oct 2026 00:53:52 +0200 Subject: [PATCH 1/2] Stream the buddy agent turn live from the AI service - Turn the agent loop into a generator that streams reasoning, tokens and search announcements per hop - Add POST /onboarding/buddy/agent/stream; the sync endpoint drains the same generator - Hold back text that a filter override could replace, and join hop texts with a blank line --- src/api/routes/buddy.py | 152 +++++++-- src/onboarding/buddy_agent.py | 209 ++++++++++-- tests/api/test_buddy.py | 101 ++++++ tests/onboarding/test_buddy_agent.py | 20 +- tests/onboarding/test_buddy_agent_stream.py | 314 +++++++++++++++++++ tests/onboarding/test_buddy_search_parity.py | 23 +- 6 files changed, 739 insertions(+), 80 deletions(-) create mode 100644 tests/onboarding/test_buddy_agent_stream.py diff --git a/src/api/routes/buddy.py b/src/api/routes/buddy.py index 0131792..6e54525 100644 --- a/src/api/routes/buddy.py +++ b/src/api/routes/buddy.py @@ -21,7 +21,14 @@ from ingestion.source_state_store import SourceStateStore from llm.base import LLMClient, Message, ToolCall, ToolSpec from llm.errors import LLMUnavailableError -from onboarding.buddy_agent import run_agent_turn +from onboarding.buddy_agent import ( + AgentReasoning, + AgentToken, + AgentToolUse, + AgentTurnResult, + run_agent_turn, + stream_agent_turn, +) from onboarding.buddy_compact import compact_memory from onboarding.buddy_open import stream_session from onboarding.vocabulary import Vocabulary @@ -67,6 +74,14 @@ def _to_toolspec(schema: BuddyToolSpecSchema) -> ToolSpec: ) +def _vocabulary(body: BuddyAgentRequest) -> Vocabulary: + return Vocabulary( + contribution_noun=body.vocabulary.contribution_noun, + contribution_noun_plural=body.vocabulary.contribution_noun_plural, + contribution_verb_past=body.vocabulary.contribution_verb_past, + ) + + def _narrowing(body: BuddyAgentRequest) -> RetrievalFilters | None: """The reader's source-system and time narrowing, if they chose any. @@ -82,6 +97,29 @@ def _narrowing(body: BuddyAgentRequest) -> RetrievalFilters | None: ) +def _agent_response(result: AgentTurnResult) -> BuddyAgentResponse: + return BuddyAgentResponse( + final=result.final, + text=result.text, + messages=[_from_message(m) for m in result.messages], + pending_tool_calls=[ + BuddyToolCallSchema( + id=call.id, name=call.name, arguments=dict(call.arguments) + ) + for call in result.pending_tool_calls + ], + citations=[ + BuddyCitationSchema( + artifact_id=cit.artifact_id, + start_line=cit.start_line, + start_page=cit.start_page, + ) + for cit in result.citations + ], + reasoning=result.reasoning, + ) + + @router.post( "/onboarding/buddy/agent", response_model=BuddyAgentResponse, @@ -104,22 +142,18 @@ def buddy_agent( ``capabilities_enabled`` and ``team_mode`` pick the persona. Both are read on every hop rather than only the first: the persona is rebuilt each time, so a resume hop that omitted one would finish the turn in the other mode. + + The same turn as ``/onboarding/buddy/agent/stream``, buffered into one response. """ - messages = [_to_message(m) for m in body.messages] - backend_tools = [_to_toolspec(t) for t in body.backend_tools] try: result = run_agent_turn( - messages, - backend_tools, + [_to_message(m) for m in body.messages], + [_to_toolspec(t) for t in body.backend_tools], llm, store, exclusions=source_state.get_exclusions(), prior_summary=body.prior_summary, - vocabulary=Vocabulary( - contribution_noun=body.vocabulary.contribution_noun, - contribution_noun_plural=body.vocabulary.contribution_noun_plural, - contribution_verb_past=body.vocabulary.contribution_verb_past, - ), + vocabulary=_vocabulary(body), project_ids=frozenset(body.project_ids), capabilities_enabled=body.capabilities_enabled, team_mode=body.team_mode, @@ -130,26 +164,86 @@ def buddy_agent( status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail=str(exc) ) from exc - return BuddyAgentResponse( - final=result.final, - text=result.text, - messages=[_from_message(m) for m in result.messages], - pending_tool_calls=[ - BuddyToolCallSchema( - id=call.id, name=call.name, arguments=dict(call.arguments) - ) - for call in result.pending_tool_calls - ], - citations=[ - BuddyCitationSchema( - artifact_id=cit.artifact_id, - start_line=cit.start_line, - start_page=cit.start_page, + return _agent_response(result) + + +@router.post( + "/onboarding/buddy/agent/stream", + summary="Run one agentic buddy turn, streaming it as it happens", + response_class=StreamingResponse, + tags=["onboarding-buddy"], + responses={422: {"model": ValidationErrorResponse}}, +) +def buddy_agent_stream( + body: BuddyAgentRequest, + llm: LLMClient = Depends(get_llm), + store: VectorStore = Depends(get_store), + source_state: SourceStateStore = Depends(get_source_state_store), +) -> StreamingResponse: + """One turn of the tool-using buddy, with the model's work visible as it happens. + + Takes the same request as ``/onboarding/buddy/agent`` and runs the same turn; + what differs is that nothing waits for the last hop. Emits ``reasoning`` and + ``token`` events, each carrying a ``content`` fragment, as the model produces + them; a ``tool_use`` event (``name``, ``arguments``) before each search this + service runs itself; and one terminal ``result`` carrying the fields of + ``BuddyAgentResponse`` (``final``, ``text``, ``messages``, ``pending_tool_calls``, + ``citations``, ``reasoning``). The ``text`` of ``result`` is authoritative: it is + what the caller stores and carries back, and the ``token`` events spell it out + except where a turn that streamed text for one hop then searched again, which + joins its hops with a blank line. + + A model that is unavailable, or any failure mid-turn, becomes a terminal ``error`` + event rather than a truncated body; the status line is already sent by then. + """ + messages = [_to_message(m) for m in body.messages] + backend_tools = [_to_toolspec(t) for t in body.backend_tools] + exclusions = source_state.get_exclusions() + narrowing = _narrowing(body) + + def event_stream() -> Iterator[str]: + try: + for event in stream_agent_turn( + messages, + backend_tools, + llm, + store, + exclusions=exclusions, + prior_summary=body.prior_summary, + vocabulary=_vocabulary(body), + project_ids=frozenset(body.project_ids), + capabilities_enabled=body.capabilities_enabled, + team_mode=body.team_mode, + filters=narrowing, + ): + if isinstance(event, AgentReasoning): + yield sse_event({"type": "reasoning", "content": event.text}) + elif isinstance(event, AgentToken): + yield sse_event({"type": "token", "content": event.text}) + elif isinstance(event, AgentToolUse): + yield sse_event( + { + "type": "tool_use", + "name": event.name, + "arguments": event.arguments, + } + ) + else: + yield sse_event( + { + "type": "result", + **_agent_response(event).model_dump(mode="json"), + } + ) + except LLMUnavailableError as exc: + yield sse_event({"type": "error", "message": str(exc)}) + except Exception: + logger.exception("Unexpected error in buddy agent stream") + yield sse_event( + {"type": "error", "message": "An unexpected error occurred"} ) - for cit in result.citations - ], - reasoning=result.reasoning, - ) + + return StreamingResponse(event_stream(), media_type="text/event-stream") @router.post( diff --git a/src/onboarding/buddy_agent.py b/src/onboarding/buddy_agent.py index cc88e3f..036f926 100644 --- a/src/onboarding/buddy_agent.py +++ b/src/onboarding/buddy_agent.py @@ -25,16 +25,28 @@ Folding older turns into that summary is ``onboarding.buddy_compact``, on its own endpoint, which the backend runs after a turn rather than during one -- see ``run_agent_turn`` for what that cost while it lived here. + +The loop is ``stream_agent_turn``, a generator, so a reader sees the model think and +write while the turn runs instead of after its last hop. ``run_agent_turn`` drains it +for the callers that only want the outcome. """ -from collections.abc import Collection +from collections.abc import Collection, Iterator from dataclasses import dataclass, field, replace from agents.tools.base import ToolRegistry, ToolResult, run_tool_calls from agents.tools.evidence import format_evidence, limit_evidence from agents.tools.grep import GrepTool from agents.tools.retrieve import RetrieveTool -from llm.base import ChatResult, LLMClient, Message, ToolCall, ToolSpec +from llm.base import ( + ChatResult, + LLMClient, + Message, + ReasoningDelta, + TextDelta, + ToolCall, + ToolSpec, +) from llm.errors import LLMUnavailableError from onboarding.buddy_persona import build_persona from onboarding.query_fence import QUERY_FENCE_NOTE, fence_user_messages @@ -129,6 +141,56 @@ class AgentTurnResult: reasoning: list[str] = field(default_factory=list[str]) +@dataclass(frozen=True) +class AgentReasoning: + """A fragment of the model's reasoning, as it arrives.""" + + text: str + + +@dataclass(frozen=True) +class AgentToken: + """A fragment of the answer text, as it arrives.""" + + text: str + + +@dataclass(frozen=True) +class AgentToolUse: + """A search this loop is about to run itself (never a backend tool).""" + + name: str + arguments: dict[str, object] + + +type AgentStreamEvent = AgentReasoning | AgentToken | AgentToolUse | AgentTurnResult + +# Between the text of one hop and the text of the next. A hop that wrote something and +# then searched leaves its text without a trailing break, and the next hop would run +# straight on from it ("...let me check.The deploy runs on..."). +_HOP_SEPARATOR = "\n\n" + + +class _HopText: + """Keeps the answer text of successive hops readable as one stream.""" + + def __init__(self) -> None: + self._emitted = False + self._owed = False + + def begin_hop(self) -> None: + self._owed = self._emitted + + def token(self, text: str) -> Iterator[AgentToken]: + if not text: + return + if self._owed: + self._owed = False + yield AgentToken(_HOP_SEPARATOR) + self._emitted = True + yield AgentToken(text) + + def _persona_prompt( summary: str | None, tool_names: Collection[str], @@ -297,7 +359,7 @@ def _local_tools( ) -def run_agent_turn( +def stream_agent_turn( messages: list[Message], backend_tools: list[ToolSpec], llm: LLMClient, @@ -309,14 +371,20 @@ def run_agent_turn( capabilities_enabled: bool = True, team_mode: bool = False, filters: RetrievalFilters | None = None, -) -> AgentTurnResult: - """Runs one agent turn: executes the searches locally, pauses on backend tools. +) -> Iterator[AgentStreamEvent]: + """Runs one agent turn, yielding what the model does as it does it. + + Executes the searches locally and pauses on backend tools. Loops internally while + the model only asks for local searches (``search_docs``, and ``grep`` when + capabilities are off); finishes as soon as it either produces a final answer or + requests a tool only the backend can run. A step budget bounds the internal loop; + if it's exhausted the model is asked once more with no tools, forcing an answer. - Loops internally while the model only asks for local searches (``search_docs``, - and ``grep`` when capabilities are off); returns as soon as it either produces a - final answer or requests a tool only the backend can run. A step budget bounds - the internal loop; if it's exhausted the model is asked once more with no tools, - forcing an answer. + Yields ``AgentReasoning`` and ``AgentToken`` fragments as the model produces them, + an ``AgentToolUse`` before each search this loop runs, and as the last event the + ``AgentTurnResult`` -- the same outcome ``run_agent_turn`` returns. Text from a + hop that went on to search is streamed like any other, and the next hop's text is + set apart from it by a blank line, written here and nowhere else. ``prior_summary`` stands in for everything older than ``messages``. @@ -328,7 +396,10 @@ def run_agent_turn( are ignored, because the scope is ``project_ids``. When it narrows, this hop searched, nothing it found survived, and no backend tool is mounted to have supplied other evidence, the answer is ``NO_FILTERED_RESULTS_MESSAGE`` rather - than whatever the model composed from nothing. + than whatever the model composed from nothing. Whether a hop can end that way is + known before it starts, so such a hop's text is held back rather than streamed: + it is released if the hop goes on to search, and replaced by the message if it is + the answer. This turn does not fold anything, and used not to be able to say that. A ``summarize_upto`` argument asked it to compact the oldest window messages @@ -378,12 +449,16 @@ def run_agent_turn( reasoning: list[str] = [] searched = False found = False + text_out = _HopText() - def _answer(text: str) -> AgentTurnResult: - # Nothing matched under an explicit narrowing: the reader chose the filter, - # so an answer composed from no sources is the one thing not to give. + # Nothing matched under an explicit narrowing: the reader chose the filter, so an + # answer composed from no sources is the one thing not to give. + def _overridable() -> bool: narrowed = has_narrowing_filters(filters) - if narrowed and searched and not found and not backend_names: + return narrowed and searched and not found and not backend_names + + def _answer(text: str) -> AgentTurnResult: + if _overridable(): text = NO_FILTERED_RESULTS_MESSAGE return AgentTurnResult( final=True, @@ -394,15 +469,42 @@ def _answer(text: str) -> AgentTurnResult: ) for _ in range(_MAX_STEPS): - result = llm.chat(work, tools) + text_out.begin_hop() + # A hop that can end as the filter message must not have shown the model's + # own text first, so it is held until the hop says which it is. + hold = _overridable() + held: list[str] = [] + result: ChatResult | None = None + for event in llm.chat_stream(work, tools): + if isinstance(event, ReasoningDelta): + if event.text: + yield AgentReasoning(event.text) + elif isinstance(event, TextDelta): + if hold: + held.append(event.text) + else: + yield from text_out.token(event.text) + else: + result = event + if result is None: + raise LLMUnavailableError("The model ended its reply without a result.") if result.reasoning and result.reasoning.strip(): reasoning.append(result.reasoning) if not result.tool_calls: - return _answer(result.text) + answer = _answer(result.text) + if hold: + yield from text_out.token(answer.text) + yield answer + return + # Not the answer after all: the held text was narration before a search. + for fragment in held: + yield from text_out.token(fragment) work = [*work, _assistant_message(result)] local_calls = [c for c in result.tool_calls if c.name in local_names] + for call in local_calls: + yield AgentToolUse(call.name, dict(call.arguments)) outcomes = dict( zip( (c.id for c in local_calls), @@ -440,7 +542,7 @@ def _answer(text: str) -> AgentTurnResult: # in this same turn already have their results appended above, so the message # list stays well-formed. if pending: - return AgentTurnResult( + yield AgentTurnResult( final=False, text=result.text, messages=work, @@ -448,23 +550,84 @@ def _answer(text: str) -> AgentTurnResult: citations=citations, reasoning=reasoning, ) + return # Only local searches this turn -- loop and let the model reason over them. # Step budget spent: force a final answer with no tools rather than loop forever. # # The notice goes to the model, not into the returned conversation: a final turn's # messages are discarded by the caller, and a transcript carrying instructions about - # a budget nobody can see would be a strange thing to keep. `generate` returns - # text only, so this call adds nothing to ``reasoning``. - return _answer( - llm.generate([*work, Message(role="system", content=_NO_TOOLS_THIS_TURN)]) - ) + # a budget nobody can see would be a strange thing to keep. + text_out.begin_hop() + hold = _overridable() + forced = "" + forced_reasoning = "" + notice = Message(role="system", content=_NO_TOOLS_THIS_TURN) + for event in llm.stream([*work, notice]): + if isinstance(event, ReasoningDelta): + if event.text: + forced_reasoning += event.text + yield AgentReasoning(event.text) + else: + forced += event.text + if not hold: + yield from text_out.token(event.text) + if forced_reasoning.strip(): + reasoning.append(forced_reasoning) + answer = _answer(forced) + if hold: + yield from text_out.token(answer.text) + yield answer + + +def run_agent_turn( + messages: list[Message], + backend_tools: list[ToolSpec], + llm: LLMClient, + store: VectorStore, + exclusions: SourceExclusions | None = None, + prior_summary: str | None = None, + vocabulary: Vocabulary = DEFAULT_VOCABULARY, + project_ids: frozenset[str] | None = None, + capabilities_enabled: bool = True, + team_mode: bool = False, + filters: RetrievalFilters | None = None, +) -> AgentTurnResult: + """Runs one agent turn to its outcome; see ``stream_agent_turn`` for the rules. + + Same turn, same arguments, with the fragments discarded: what the buffered + endpoint and the tests want is the ``AgentTurnResult``. + """ + outcome: AgentTurnResult | None = None + for event in stream_agent_turn( + messages, + backend_tools, + llm, + store, + exclusions=exclusions, + prior_summary=prior_summary, + vocabulary=vocabulary, + project_ids=project_ids, + capabilities_enabled=capabilities_enabled, + team_mode=team_mode, + filters=filters, + ): + if isinstance(event, AgentTurnResult): + outcome = event + if outcome is None: + raise LLMUnavailableError("The agent turn ended without an outcome.") + return outcome __all__ = [ + "AgentReasoning", + "AgentStreamEvent", + "AgentToken", + "AgentToolUse", "AgentTurnResult", "GREP", "LLMUnavailableError", "run_agent_turn", "SEARCH_DOCS", + "stream_agent_turn", ] diff --git a/tests/api/test_buddy.py b/tests/api/test_buddy.py index b2f319f..8d96ecf 100644 --- a/tests/api/test_buddy.py +++ b/tests/api/test_buddy.py @@ -12,6 +12,7 @@ from llm.errors import LLMUnavailableError from rag.filters import NO_FILTERED_RESULTS_MESSAGE from rag.types import Chunk +from tests.conftest import parse_sse_events from tests.stubs.llm import ScriptedLLMClient, StubLLMClient from tests.stubs.store import StubVectorStore @@ -502,3 +503,103 @@ def test_a_caller_that_sends_no_filters_gets_todays_turn(client: TestClient) -> assert response.status_code == 200 assert response.json()["reasoning"] == [] + + +# --- streaming agent turn ----------------------------------------------------- + +_STREAM_URL = "/api/v1/onboarding/buddy/agent/stream" + + +def test_the_stream_endpoint_emits_progress_then_the_buffered_response() -> None: + llm = ScriptedLLMClient( + turns=[[("grep", {"patterns": ["login"]})]], answer="In auth.py." + ) + payload = { + "messages": [{"role": "user", "content": "where is login?"}], + "capabilities_enabled": False, + "project_ids": ["p1"], + } + for client in _scripted_client(llm): + streamed = client.post(_STREAM_URL, json=payload) + llm_again = ScriptedLLMClient( + turns=[[("grep", {"patterns": ["login"]})]], answer="In auth.py." + ) + for client in _scripted_client(llm_again): + buffered = client.post(_URL, json=payload) + + assert streamed.status_code == 200 + assert streamed.headers["content-type"].startswith("text/event-stream") + events = parse_sse_events(streamed.text) + kinds = [e["type"] for e in events] + assert kinds == ["tool_use", "token", "result"] + assert events[0]["name"] == "grep" + assert events[0]["arguments"] == {"patterns": ["login"]} + assert events[1]["content"] == "In auth.py." + result = {k: v for k, v in events[-1].items() if k != "type"} + # The terminal event is the buffered endpoint's response, field for field. + assert result == buffered.json() + assert result["citations"] == [ + {"artifact_id": "a1", "start_line": None, "start_page": None} + ] + + +def test_the_stream_endpoint_emits_reasoning_deltas() -> None: + llm = ScriptedLLMClient( + turns=[[("grep", {"patterns": ["login"]})]], reasoning="grep for it" + ) + for client in _scripted_client(llm): + response = client.post( + _STREAM_URL, + json={ + "messages": [{"role": "user", "content": "login?"}], + "capabilities_enabled": False, + }, + ) + + events = parse_sse_events(response.text) + assert {"type": "reasoning", "content": "grep for it"} in events + assert events[-1]["reasoning"] == ["grep for it"] + + +def test_the_stream_endpoint_reports_an_unavailable_model_as_an_error_event() -> None: + class _Down(StubLLMClient): + def chat_stream(self, *args: Any, **kwargs: Any) -> Any: + raise LLMUnavailableError("model down") + + app.dependency_overrides[get_llm] = lambda: _Down() + app.dependency_overrides[get_store] = lambda: StubVectorStore() + app.dependency_overrides[get_source_state_store] = lambda: SourceStateStore( + ":memory:" + ) + try: + response = TestClient(app).post( + _STREAM_URL, json={"messages": [{"role": "user", "content": "hi"}]} + ) + finally: + app.dependency_overrides.clear() + + assert response.status_code == 200 + events = parse_sse_events(response.text) + assert [e["type"] for e in events] == ["error"] + assert "model down" in events[0]["message"] + + +def test_the_stream_endpoint_applies_the_filters_to_its_answer() -> None: + llm = ScriptedLLMClient( + turns=[[("grep", {"patterns": ["login handler"]})]], answer="ungrounded" + ) + for client in _scripted_client(llm): + response = client.post( + _STREAM_URL, + json={ + "messages": [{"role": "user", "content": "where is login?"}], + "capabilities_enabled": False, + "filters": {"source_systems": ["jira"]}, + "project_ids": ["p1"], + }, + ) + + events = parse_sse_events(response.text) + streamed = "".join(e["content"] for e in events if e["type"] == "token") + assert streamed == NO_FILTERED_RESULTS_MESSAGE + assert events[-1]["text"] == NO_FILTERED_RESULTS_MESSAGE diff --git a/tests/onboarding/test_buddy_agent.py b/tests/onboarding/test_buddy_agent.py index 23fa581..0126f68 100644 --- a/tests/onboarding/test_buddy_agent.py +++ b/tests/onboarding/test_buddy_agent.py @@ -231,20 +231,6 @@ def _fake_retrieve(*args: object, **kwargs: object) -> list[ScoredChunk]: assert [getattr(f, "project_ids", None) for f in seen] == [None] -class _RecordingForcedAnswer(ScriptedLLMClient): - """Records what the forced, tool-less final answer was asked with.""" - - def __init__(self, *args: object, **kwargs: object) -> None: - super().__init__(*args, **kwargs) # type: ignore[arg-type] - self.generate_calls: list[list[Message]] = [] - - def generate( - self, messages: list[Message], *, temperature: float | None = None - ) -> str: - self.generate_calls.append(messages) - return self.answer - - def test_a_spent_search_budget_tells_the_model_it_has_no_tools( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -252,15 +238,15 @@ def test_a_spent_search_budget_tells_the_model_it_has_no_tools( # with the persona still saying "offer X", and told the hire to confirm a button # no tool call had produced. monkeypatch.setattr("agents.tools.retrieve.retrieve", lambda *a, **k: []) - llm = _RecordingForcedAnswer( + llm = ScriptedLLMClient( turns=[[(SEARCH_DOCS, {"query": "again"})]] * 20, answer="Here is what I know." ) result = run_agent_turn([_user("hi")], [_GET_MY_METRICS], llm, StubVectorStore()) assert result.final is True - assert len(llm.generate_calls) == 1 - notice = llm.generate_calls[0][-1] + assert len(llm.stream_calls) == 1 + notice = llm.stream_calls[0][-1] assert notice["role"] == "system" assert "no confirm button can appear" in (notice.get("content") or "") # The notice is for the model only; it is not kept in the conversation. diff --git a/tests/onboarding/test_buddy_agent_stream.py b/tests/onboarding/test_buddy_agent_stream.py new file mode 100644 index 0000000..277f3bf --- /dev/null +++ b/tests/onboarding/test_buddy_agent_stream.py @@ -0,0 +1,314 @@ +"""The buddy turn as a stream: what a reader sees while the loop is still running.""" + +from collections.abc import Iterator, Sequence + +import pytest + +from llm.base import ( + ChatResult, + LLMChatStreamEvent, + LLMStreamEvent, + Message, + ReasoningDelta, + TextDelta, + ToolCall, + ToolSpec, +) +from llm.errors import LLMUnavailableError +from onboarding.buddy_agent import ( + GREP, + AgentReasoning, + AgentStreamEvent, + AgentToken, + AgentToolUse, + AgentTurnResult, + run_agent_turn, + stream_agent_turn, +) +from rag.filters import NO_FILTERED_RESULTS_MESSAGE +from rag.types import Chunk, RetrievalFilters +from tests.stubs.llm import ScriptedLLMClient +from tests.stubs.store import StubVectorStore + +_EMBEDDING = [1.0] + [0.0] * 767 +_METRICS: ToolSpec = { + "name": "get_my_metrics", + "description": "The hire's onboarding metrics.", + "parameters": {"type": "object", "properties": {}}, +} +_JIRA_ONLY = RetrievalFilters(source_systems=["JIRA"]) + +type Hop = Sequence[LLMChatStreamEvent] + + +def _user(text: str) -> Message: + return Message(role="user", content=text) + + +def _store() -> StubVectorStore: + store = StubVectorStore() + store.add( + [ + Chunk( + id="c1", + artifact_id="a1", + filename="auth.md", + text="the login handler lives in auth.py", + embedding=_EMBEDDING, + source_system="GITHUB", + ) + ] + ) + return store + + +def _search(query: str = "login") -> ChatResult: + return ChatResult( + text="", + tool_calls=[ToolCall(id="call_0", name=GREP, arguments={"patterns": [query]})], + ) + + +def _narrated_search(text: str, query: str = "login") -> Hop: + """A hop that says something, then asks for a search.""" + call = ToolCall(id="call_0", name=GREP, arguments={"patterns": [query]}) + return [TextDelta(text), ChatResult(text=text, tool_calls=[call])] + + +def _answer(*fragments: str) -> Hop: + return [*(TextDelta(f) for f in fragments), ChatResult(text="".join(fragments))] + + +class _Hops(ScriptedLLMClient): + """Plays one scripted event sequence per model call.""" + + def __init__(self, hops: Sequence[Hop], forced: Sequence[LLMStreamEvent] = ()): + super().__init__(turns=[]) + self._hops = list(hops) + self._forced = list(forced) + + def chat_stream( + self, messages: list[Message], tools: list[ToolSpec] | None = None + ) -> Iterator[LLMChatStreamEvent]: + self.chat_calls.append(messages) + yield from self._hops.pop(0) + + def stream(self, messages: list[Message]) -> Iterator[LLMStreamEvent]: + self.stream_calls.append(messages) + yield from self._forced + + +def _events( + llm: ScriptedLLMClient, + *, + backend_tools: list[ToolSpec] | None = None, + filters: RetrievalFilters | None = None, +) -> list[AgentStreamEvent]: + return list( + stream_agent_turn( + [_user("where is login?")], + backend_tools or [], + llm, + _store(), + capabilities_enabled=False, + filters=filters, + ) + ) + + +def _text(events: list[AgentStreamEvent]) -> str: + return "".join(e.text for e in events if isinstance(e, AgentToken)) + + +def _outcome(events: list[AgentStreamEvent]) -> AgentTurnResult: + last = events[-1] + assert isinstance(last, AgentTurnResult) + return last + + +def test_reasoning_and_tokens_arrive_in_order_and_the_result_comes_last() -> None: + llm = _Hops( + [ + [ + ReasoningDelta("Let me "), + ReasoningDelta("think."), + TextDelta("In "), + TextDelta("auth.py."), + ChatResult(text="In auth.py.", reasoning="Let me think."), + ] + ] + ) + + events = _events(llm) + + assert events[:-1] == [ + AgentReasoning("Let me "), + AgentReasoning("think."), + AgentToken("In "), + AgentToken("auth.py."), + ] + outcome = _outcome(events) + assert outcome.final is True + assert outcome.text == "In auth.py." + assert outcome.reasoning == ["Let me think."] + + +def test_fragments_are_yielded_before_the_model_has_finished() -> None: + state = {"finished": False} + + class _Slow(ScriptedLLMClient): + def chat_stream( + self, messages: list[Message], tools: list[ToolSpec] | None = None + ) -> Iterator[LLMChatStreamEvent]: + yield TextDelta("first") + state["finished"] = True + yield ChatResult(text="first") + + turn = stream_agent_turn([_user("hi")], [], _Slow(turns=[]), _store()) + + assert next(turn) == AgentToken("first") + assert state["finished"] is False + + +def test_each_local_search_is_announced_before_it_runs() -> None: + llm = _Hops([[_search("login")], _answer("ok")]) + + events = _events(llm) + + uses = [e for e in events if isinstance(e, AgentToolUse)] + assert uses == [AgentToolUse(GREP, {"patterns": ["login"]})] + assert events.index(uses[0]) < events.index(_outcome(events)) + + +def test_a_backend_tool_is_not_announced_as_a_search() -> None: + call = ToolCall(id="m1", name="get_my_metrics", arguments={}) + llm = _Hops([[ChatResult(text="", tool_calls=[call])]]) + + events = _events(llm, backend_tools=[_METRICS]) + + assert not any(isinstance(e, AgentToolUse) for e in events) + outcome = _outcome(events) + assert outcome.final is False + assert [c.name for c in outcome.pending_tool_calls] == ["get_my_metrics"] + + +def test_the_text_of_one_hop_is_set_apart_from_the_next() -> None: + llm = _Hops([_narrated_search("Let me check."), _answer("It is in auth.py.")]) + + events = _events(llm) + + assert _text(events) == "Let me check.\n\nIt is in auth.py." + # Only the answer is the stored text; the narration was streamed, not kept. + assert _outcome(events).text == "It is in auth.py." + + +def test_a_hop_that_wrote_nothing_adds_no_separator() -> None: + llm = _Hops([[_search()], _answer("It is in auth.py.")]) + + assert _text(_events(llm)) == "It is in auth.py." + + +def test_the_separator_survives_a_silent_hop_in_between() -> None: + llm = _Hops( + [_narrated_search("Checking."), [_search("again")], _answer("Found it.")] + ) + + assert _text(_events(llm)) == "Checking.\n\nFound it." + + +def test_a_filtered_answer_with_no_sources_is_replaced_not_streamed() -> None: + llm = _Hops([[_search()], _answer("an ", "ungrounded ", "guess")]) + + events = _events(llm, filters=_JIRA_ONLY) + + assert _text(events) == NO_FILTERED_RESULTS_MESSAGE + assert _outcome(events).text == NO_FILTERED_RESULTS_MESSAGE + + +def test_held_text_is_released_when_the_hop_goes_on_to_search() -> None: + llm = _Hops( + [ + [_search()], + _narrated_search("Still looking.", "other"), + _answer("guess"), + ] + ) + + events = _events(llm, filters=_JIRA_ONLY) + + assert _text(events) == f"Still looking.\n\n{NO_FILTERED_RESULTS_MESSAGE}" + + +def test_a_filtered_search_that_found_something_streams_normally() -> None: + llm = _Hops([[_search()], _answer("In ", "auth.py.")]) + + events = _events(llm, filters=RetrievalFilters(source_systems=["GITHUB"])) + + assert _text(events) == "In auth.py." + assert _outcome(events).text == "In auth.py." + + +def test_a_spent_search_budget_streams_the_forced_answer() -> None: + llm = _Hops( + [[_search()]] * 6, + forced=[ + ReasoningDelta("wrapping up"), + TextDelta("Here is "), + TextDelta("all."), + ], + ) + + events = _events(llm) + + assert len(llm.stream_calls) == 1 + assert AgentReasoning("wrapping up") in events + assert _text(events) == "Here is all." + outcome = _outcome(events) + assert outcome.final is True + assert outcome.text == "Here is all." + assert outcome.reasoning == ["wrapping up"] + + +def test_a_forced_answer_under_a_filter_with_no_sources_is_replaced() -> None: + llm = _Hops([[_search()]] * 6, forced=[TextDelta("guess")]) + + events = _events(llm, filters=_JIRA_ONLY) + + assert _text(events) == NO_FILTERED_RESULTS_MESSAGE + + +def test_abandoning_the_stream_stops_the_loop() -> None: + """The caller hanging up mid-turn must not leave the model running more hops.""" + llm = _Hops([_narrated_search("Let me check."), _answer("never reached")]) + turn = stream_agent_turn( + [_user("hi")], [], llm, _store(), capabilities_enabled=False + ) + + assert next(turn) == AgentToken("Let me check.") + turn.close() + + assert len(llm.chat_calls) == 1 + + +def test_a_model_that_ends_without_a_result_is_an_unavailable_model() -> None: + llm = _Hops([[TextDelta("half an ans")]]) + + with pytest.raises(LLMUnavailableError): + _events(llm) + + +def test_the_buffered_turn_is_the_streamed_turns_outcome() -> None: + def make() -> _Hops: + return _Hops([_narrated_search("Checking."), _answer("In auth.py.")]) + + streamed = _outcome(_events(make())) + buffered = run_agent_turn( + [_user("where is login?")], + [], + make(), + _store(), + capabilities_enabled=False, + ) + + assert buffered == streamed + assert buffered.messages[-1] == Message(role="assistant", content="In auth.py.") diff --git a/tests/onboarding/test_buddy_search_parity.py b/tests/onboarding/test_buddy_search_parity.py index ce5c823..aa538cb 100644 --- a/tests/onboarding/test_buddy_search_parity.py +++ b/tests/onboarding/test_buddy_search_parity.py @@ -5,11 +5,11 @@ moving to the buddy must not lose. """ -from collections.abc import Sequence +from collections.abc import Iterator, Sequence import pytest -from llm.base import ChatResult, Message, ToolSpec +from llm.base import ChatResult, LLMChatStreamEvent, Message, ToolSpec from onboarding.buddy_agent import ( GREP, SEARCH_DOCS, @@ -43,11 +43,11 @@ def __init__( super().__init__(turns, answer=answer, reasoning=reasoning) self.offered: list[list[str]] = [] - def chat( + def chat_stream( self, messages: list[Message], tools: list[ToolSpec] | None = None - ) -> ChatResult: + ) -> Iterator[LLMChatStreamEvent]: self.offered.append([tool["name"] for tool in tools or []]) - return super().chat(messages, tools) + return super().chat_stream(messages, tools) def _user(text: str) -> Message: @@ -347,13 +347,14 @@ def test_several_searches_in_one_step_all_run_and_answer_in_call_order() -> None class _ReasonsEveryStep(ScriptedLLMClient): """Returns reasoning on the final step too, which the base stub does not.""" - def chat( + def chat_stream( self, messages: list[Message], tools: list[ToolSpec] | None = None - ) -> ChatResult: - result = super().chat(messages, tools) - if result.tool_calls: - return result - return ChatResult(text=result.text, reasoning="Now I can answer.") + ) -> Iterator[LLMChatStreamEvent]: + for event in super().chat_stream(messages, tools): + if isinstance(event, ChatResult) and not event.tool_calls: + yield ChatResult(text=event.text, reasoning="Now I can answer.") + else: + yield event def test_reasoning_is_returned_per_step_and_kept_out_of_the_transcript() -> None: From 9cb672a9a350f97127b90a948267b02e3d356cf7 Mon Sep 17 00:00:00 2001 From: kiranfin <4finjas@online.de> Date: Wed, 7 Oct 2026 02:24:11 +0200 Subject: [PATCH 2/2] Carry the buddy's reasoning through backend tool round trips - Add reasoning and reasoning_details to the buddy agent wire message and map them both ways - Fix the resume hop after a backend tool, which a thinking-enabled provider rejected with a 400 - Replace the test that pinned reasoning out of messages with carry and resume tests --- src/api/routes/buddy.py | 9 +++++ src/api/schemas.py | 16 +++++++++ src/onboarding/buddy_agent.py | 7 ++-- tests/api/test_buddy.py | 63 ++++++++++++++++++++++++++++------- 4 files changed, 80 insertions(+), 15 deletions(-) diff --git a/src/api/routes/buddy.py b/src/api/routes/buddy.py index 6e54525..0fab3ae 100644 --- a/src/api/routes/buddy.py +++ b/src/api/routes/buddy.py @@ -49,10 +49,17 @@ def _to_message(schema: BuddyAgentMessageSchema) -> Message: ] if schema.tool_call_id is not None: msg["tool_call_id"] = schema.tool_call_id + if schema.reasoning: + msg["reasoning"] = schema.reasoning + if schema.reasoning_details: + msg["reasoning_details"] = [dict(d) for d in schema.reasoning_details] return msg def _from_message(msg: Message) -> BuddyAgentMessageSchema: + # The reasoning rides along so the hop after a backend tool can hand it back: + # a provider with extended thinking rejects a tool turn that lost its signed + # thinking blocks. return BuddyAgentMessageSchema( role=msg["role"], content=msg.get("content") or "", @@ -63,6 +70,8 @@ def _from_message(msg: Message) -> BuddyAgentMessageSchema: for call in msg.get("tool_calls") or [] ], tool_call_id=msg.get("tool_call_id"), + reasoning=msg.get("reasoning"), + reasoning_details=list(msg.get("reasoning_details") or []), ) diff --git a/src/api/schemas.py b/src/api/schemas.py index 1fea272..f7a787e 100644 --- a/src/api/schemas.py +++ b/src/api/schemas.py @@ -1303,6 +1303,22 @@ class BuddyAgentMessageSchema(BaseModel): default=None, description="On a tool-result turn, the id of the call it answers.", ) + reasoning: str | None = Field( + default=None, + description=( + "On an assistant turn, the model's plain-text reasoning. Opaque to the " + "caller; carry it back verbatim." + ), + ) + reasoning_details: list[dict[str, object]] = Field( + default_factory=list[dict[str, object]], + description=( + "On an assistant turn, the provider's structured (possibly signed) " + "reasoning blocks. Opaque to the caller; carry them back verbatim, " + "because a provider with extended thinking rejects a resumed tool turn " + "that lost them." + ), + ) class BuddyToolSpecSchema(BaseModel): diff --git a/src/onboarding/buddy_agent.py b/src/onboarding/buddy_agent.py index 036f926..134256a 100644 --- a/src/onboarding/buddy_agent.py +++ b/src/onboarding/buddy_agent.py @@ -252,9 +252,10 @@ def _assistant_message(result: ChatResult) -> Message: msg = Message(role="assistant", content=result.text) if result.tool_calls: msg["tool_calls"] = result.tool_calls - # Kept for the next internal hop, as the chat agent does: a reasoning provider - # continues a tool-using thought from these. They never reach the caller -- - # the wire message has no field for them (``api.routes.buddy._from_message``). + # Kept for the next hop: a reasoning provider continues a tool-using thought + # from these, and one with extended thinking rejects the turn without them. + # They also go out on the wire message (``api.routes.buddy._from_message``), + # so the hop after a backend tool gets them back from the caller. if result.reasoning: msg["reasoning"] = result.reasoning if result.reasoning_details: diff --git a/tests/api/test_buddy.py b/tests/api/test_buddy.py index 8d96ecf..d7f243d 100644 --- a/tests/api/test_buddy.py +++ b/tests/api/test_buddy.py @@ -472,28 +472,67 @@ def test_an_empty_project_scope_admits_nothing( assert "retro.md" not in seen -def test_reasoning_is_returned_and_never_carried_in_messages() -> None: +_SIGNED_DETAILS: list[dict[str, Any]] = [ + {"type": "reasoning.text", "text": "check the board", "signature": "sig-1"} +] + +_READ_BOARD_SPEC = { + "name": "read_board", + "description": "Reads the hire's board.", + "parameters": {"type": "object", "properties": {}}, +} + + +def test_a_turn_waiting_on_a_backend_tool_carries_its_reasoning() -> None: llm = ScriptedLLMClient( - turns=[[("grep", {"patterns": ["login"]})]], - reasoning="grep for it", - reasoning_details=[{"type": "reasoning.text", "text": "grep for it"}], + turns=[[("read_board", {})]], + reasoning="check the board", + reasoning_details=_SIGNED_DETAILS, ) for client in _scripted_client(llm): response = client.post( _URL, json={ - "messages": [{"role": "user", "content": "login?"}], - "capabilities_enabled": False, + "messages": [{"role": "user", "content": "what is next?"}], + "backend_tools": [_READ_BOARD_SPEC], }, ) body = response.json() - assert body["reasoning"] == ["grep for it"] - assert all( - set(m) <= {"role", "content", "tool_calls", "tool_call_id"} - for m in body["messages"] - ) - assert "grep for it" not in str(body["messages"]) + assert body["final"] is False + assert body["reasoning"] == ["check the board"] + waiting = next(m for m in body["messages"] if m["tool_calls"]) + assert waiting["reasoning"] == "check the board" + assert waiting["reasoning_details"] == _SIGNED_DETAILS + + +def test_a_resumed_turn_hands_the_carried_reasoning_back_to_the_model() -> None: + llm = ScriptedLLMClient(turns=[], answer="Your next step is the setup.") + for client in _scripted_client(llm): + response = client.post( + _URL, + json={ + "messages": [ + {"role": "user", "content": "what is next?"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + {"id": "call_0", "name": "read_board", "arguments": {}} + ], + "reasoning": "check the board", + "reasoning_details": _SIGNED_DETAILS, + }, + {"role": "tool", "content": "setup", "tool_call_id": "call_0"}, + ], + "backend_tools": [_READ_BOARD_SPEC], + }, + ) + + assert response.status_code == 200 + sent = next(m for m in llm.chat_calls[0] if m.get("tool_calls")) + assert sent.get("reasoning_details") == _SIGNED_DETAILS + assert sent.get("reasoning") == "check the board" def test_a_caller_that_sends_no_filters_gets_todays_turn(client: TestClient) -> None: