27 m45 extended thinking sur investigator opus 47 wow factor

backend/agents/investigator.py

@@ -18,11 +18,14 @@
   in ``status='analyzed'`` with a ``rca_summary`` explaining why — the
   operator always sees an outcome, the pipeline never hangs.
 
-Extended thinking (#27) is intentionally NOT enabled here yet — that
-issue flips on ``thinking={"type": "enabled", ...}`` and starts
-streaming ``thinking_delta`` events. The assistant-turn append below
-already preserves every content block verbatim, so the signed-thinking-
-block contract is satisfied the moment thinking is turned on.
+Extended thinking (#27) is enabled here via :func:`_llm_call` which wraps
+``anthropic.messages.stream()`` with ``thinking={"type": "enabled",
+"budget_tokens": 10000}`` and broadcasts each ``thinking_delta`` chunk
+as a ``thinking_delta`` WebSocket frame. The frontend Agent Inspector
+(M8.5 / #49) renders the live reasoning trace. Signed thinking blocks
+are preserved verbatim by ``response.content[*].model_dump()`` so the
+next API call (the one carrying ``tool_result`` blocks) does not trip
+the ``thinking block signature invalid`` 400 error.
 """
 
 from __future__ import annotations
@@ -36,7 +39,7 @@
 
 from agents.anthropic_client import anthropic, model_for
 from agents.ui_tools import INVESTIGATOR_RENDER_TOOLS
-from anthropic.types import ToolUseBlock
+from anthropic.types import Message, ToolUseBlock
 from aria_mcp.client import mcp_client
 from core.database import db
 from core.ws_manager import current_turn_id, ws_manager
@@ -47,7 +50,11 @@
 
 MAX_TURNS = 12
 _TIMEOUT_SECONDS = 120.0
-_MAX_TOKENS = 4096
+# Total output budget. Anthropic requires ``max_tokens > thinking.budget_tokens``
+# — keep at least 4096 tokens of headroom above the thinking budget for the
+# actual text/tool_use output.
+_THINKING_BUDGET = 10000
+_MAX_TOKENS = 16384
 
 
 # ---------------------------------------------------------------------------
@@ -151,6 +158,62 @@
 """
 
 
+# ---------------------------------------------------------------------------
+# LLM streaming + extended thinking (#27)
+# ---------------------------------------------------------------------------
+
+
+async def _llm_call(
+    *,
+    system: str,
+    messages: list[dict[str, Any]],
+    tools: list[dict[str, Any]],
+    turn_id: str,
+) -> Message:
+    """One streamed Investigator turn with extended thinking enabled.
+
+    Wraps ``anthropic.messages.stream(...)`` with
+    ``thinking={"type": "enabled", "budget_tokens": _THINKING_BUDGET}`` and
+    fans out each ``thinking_delta`` chunk as a ``thinking_delta`` WebSocket
+    frame matching ``EventBusMap.thinking_delta`` in
+    ``frontend/src/lib/ws.types.ts`` (``{agent, content, turn_id}``).
+
+    The reconstructed final ``Message`` is returned with all content blocks
+    intact — including signed ``thinking`` blocks — so the next turn's
+    ``messages.append({"role": "assistant", "content": ...})`` preserves
+    signatures and avoids the ``thinking block signature invalid`` 400.
+    """
+    async with anthropic.messages.stream(
+        model=model_for("reasoning"),
+        thinking={"type": "enabled", "budget_tokens": _THINKING_BUDGET},
+        system=system,
+        messages=cast(Any, messages),
+        tools=cast(Any, tools),
+        max_tokens=_MAX_TOKENS,
+    ) as stream:
+        async for raw_event in stream:
+            # The MessageStreamEvent union has many variants; pyright would
+            # narrow too aggressively here. Cast to Any since the runtime
+            # check (``getattr`` with default) is the safe path anyway.
+            event = cast(Any, raw_event)
+            if (
+                getattr(event, "type", None) == "content_block_delta"
+                and getattr(getattr(event, "delta", None), "type", None) == "thinking_delta"
+            ):
+                # Anthropic SDK exposes the chunk text on ``.thinking``.
+                chunk = getattr(event.delta, "thinking", None)
+                if chunk:
+                    await ws_manager.broadcast(
+                        "thinking_delta",
+                        {
+                            "agent": "investigator",
+                            "content": chunk,
+                            "turn_id": turn_id,
+                        },
+                    )
+        return await stream.get_final_message()
+
+
 # ---------------------------------------------------------------------------
 # Entry point
 # ---------------------------------------------------------------------------
@@ -225,12 +288,11 @@ async def _run_investigator_body(work_order_id: int, turn_id: str) -> None:
 
     finish_reason = "max_turns"
     for _turn in range(MAX_TURNS):
-        response = await anthropic.messages.create(
-            model=model_for("reasoning"),
+        response = await _llm_call(
             system=system_prompt,
-            messages=cast(Any, messages),
-            tools=cast(Any, tools_schema),
-            max_tokens=_MAX_TOKENS,
+            messages=messages,
+            tools=tools_schema,
+            turn_id=turn_id,
         )
         # Preserve the full assistant content verbatim — required for
         # signed `thinking` blocks the moment #27 enables thinking. Safe

backend/tests/unit/agents/test_investigator.py

@@ -31,7 +31,6 @@
 import pytest
 from agents import investigator as inv
 
-
 # ---------------------------------------------------------------------------
 # Lightweight stand-ins for the Anthropic response shape.
 # ---------------------------------------------------------------------------
@@ -70,31 +69,28 @@ class _FakeMessage:
 
 
 class _FakeAnthropic:
-    """Queue-backed fake for ``anthropic.messages.create``.
+    """Queue-backed fake exposing the same surface ``investigator._llm_call`` does.
+
+    Each invocation pops the next planned message off ``self.responses`` and
+    captures the kwargs in ``self.calls`` so tests can assert on what the
+    agent loop sent (``messages``, ``tools``, etc.).
 
-    Each call pops the next planned response off ``self.responses``.
-    The test sets up a list so the agent loop can be driven through
-    multiple turns.
+    The shape is intentionally the same as the previous ``messages.create``
+    fake so the existing test assertions (``antr.calls[N]["messages"]``)
+    keep working after the M4.5 streaming refactor.
     """
 
     def __init__(self, responses: list[_FakeMessage]) -> None:
         self.responses = list(responses)
         self.calls: list[dict[str, Any]] = []
-        self.messages = self._Messages(self)
-
-    class _Messages:
-        def __init__(self, outer: "_FakeAnthropic") -> None:
-            self._outer = outer
 
-        async def create(self, **kwargs: Any) -> _FakeMessage:
-            # Deep-copy so later ``messages.append(...)`` mutations by the
-            # agent loop do not rewrite history inside captured calls.
-            import copy
+    async def __call__(self, **kwargs: Any) -> _FakeMessage:
+        import copy
 
-            self._outer.calls.append(copy.deepcopy(kwargs))
-            if not self._outer.responses:
-                raise AssertionError("No planned Anthropic response left")
-            return self._outer.responses.pop(0)
+        self.calls.append(copy.deepcopy(kwargs))
+        if not self.responses:
+            raise AssertionError("No planned LLM response left")
+        return self.responses.pop(0)
 
 
 @dataclass
@@ -194,7 +190,11 @@ def _install(
         antr = _FakeAnthropic(responses=responses)
         mcp = _FakeMCP(results=mcp_results)
         ws = _FakeWS()
-        monkeypatch.setattr(inv, "anthropic", antr)
+        # M4.5 (#27): the loop now goes through ``_llm_call`` which wraps
+        # ``anthropic.messages.stream(...)`` with extended thinking enabled.
+        # Patch the helper directly so tests stay decoupled from the SDK
+        # streaming surface.
+        monkeypatch.setattr(inv, "_llm_call", antr)
         monkeypatch.setattr(inv, "mcp_client", mcp)
         monkeypatch.setattr(inv, "ws_manager", ws)
         monkeypatch.setattr(inv, "db", _FakeDB())
@@ -318,9 +318,7 @@ async def test_tool_call_events_carry_expected_fields(patch_inv) -> None:
 async def test_is_error_tool_result_forwarded_not_raised(patch_inv) -> None:
     """An MCP tool returning is_error=True becomes a tool_result with
     is_error=True — the loop keeps going, the LLM can self-correct."""
-    broken_call = _FakeToolUseBlock(
-        id="tu_1", name="get_signal_anomalies", input={"cell_id": 2}
-    )
+    broken_call = _FakeToolUseBlock(id="tu_1", name="get_signal_anomalies", input={"cell_id": 2})
     submit = _FakeToolUseBlock(
         id="tu_2",
         name="submit_rca",
@@ -525,9 +523,7 @@ def test_work_order_generator_lazy_import_logs_when_missing(
     monkeypatch.setitem(sys.modules, "agents.work_order_generator", None)
     with caplog.at_level("INFO"):
         inv._spawn_work_order_generator(work_order_id=42)
-    assert any(
-        "Work Order Generator not yet implemented" in rec.message for rec in caplog.records
-    )
+    assert any("Work Order Generator not yet implemented" in rec.message for rec in caplog.records)
 
 
 # ---------------------------------------------------------------------------
@@ -550,3 +546,175 @@ def test_ask_kb_builder_tool_shape() -> None:
     assert inv.ASK_KB_BUILDER_TOOL["name"] == "ask_kb_builder"
     required = inv.ASK_KB_BUILDER_TOOL["input_schema"]["required"]
     assert {"question", "cell_id"} <= set(required)
+
+
+# ---------------------------------------------------------------------------
+# M4.5 (#27) — extended thinking + thinking_delta streaming
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class _FakeStreamDelta:
+    type: str
+    thinking: str | None = None
+
+
+@dataclass
+class _FakeStreamEvent:
+    type: str
+    delta: _FakeStreamDelta | None = None
+
+
+class _FakeMessagesStream:
+    """Minimal async context manager mimicking ``anthropic.messages.stream``.
+
+    Yields the planned ``MessageStreamEvent``-like objects from ``events``
+    and returns ``final_message`` from ``get_final_message()``.
+    """
+
+    def __init__(self, events: list[_FakeStreamEvent], final_message: _FakeMessage) -> None:
+        self._events = events
+        self._final = final_message
+        self.kwargs: dict[str, Any] = {}
+
+    def __call__(self, **kwargs: Any) -> "_FakeMessagesStream":
+        self.kwargs = kwargs
+        return self
+
+    async def __aenter__(self) -> "_FakeMessagesStream":
+        return self
+
+    async def __aexit__(self, *exc: Any) -> None:
+        return None
+
+    def __aiter__(self) -> "_FakeMessagesStream":
+        self._iter = iter(self._events)
+        return self
+
+    async def __anext__(self) -> _FakeStreamEvent:
+        try:
+            return next(self._iter)
+        except StopIteration as e:
+            raise StopAsyncIteration from e
+
+    async def get_final_message(self) -> _FakeMessage:
+        return self._final
+
+
+@pytest.mark.asyncio
+async def test_llm_call_enables_thinking_and_streams_thinking_delta(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Acceptance #1 — `thinking_delta` events streamed during a turn.
+
+    Verifies that ``_llm_call`` enables extended thinking, fans out each
+    thinking chunk through ``ws_manager.broadcast`` with the
+    ``EventBusMap.thinking_delta`` shape, and returns the reconstructed
+    final message verbatim.
+    """
+    submit_block = _FakeToolUseBlock(
+        id="tu_x",
+        name="submit_rca",
+        input={
+            "root_cause": "x",
+            "failure_mode": "x",
+            "confidence": 0.5,
+            "contributing_factors": [],
+            "recommended_action": "x",
+        },
+    )
+    final = _FakeMessage(content=[submit_block], stop_reason="tool_use")
+    events = [
+        _FakeStreamEvent(type="message_start"),
+        _FakeStreamEvent(type="content_block_start"),
+        _FakeStreamEvent(
+            type="content_block_delta",
+            delta=_FakeStreamDelta(type="thinking_delta", thinking="The vibration "),
+        ),
+        _FakeStreamEvent(
+            type="content_block_delta",
+            delta=_FakeStreamDelta(type="thinking_delta", thinking="peak suggests bearing wear."),
+        ),
+        # Non-thinking deltas must NOT trigger a broadcast.
+        _FakeStreamEvent(
+            type="content_block_delta",
+            delta=_FakeStreamDelta(type="text_delta", thinking=None),
+        ),
+        _FakeStreamEvent(type="content_block_stop"),
+        _FakeStreamEvent(type="message_stop"),
+    ]
+    stream = _FakeMessagesStream(events=events, final_message=final)
+
+    class _AntStub:
+        class messages:
+            pass
+
+    _AntStub.messages.stream = stream  # type: ignore[attr-defined]
+
+    ws = _FakeWS()
+    monkeypatch.setattr(inv, "anthropic", _AntStub)
+    monkeypatch.setattr(inv, "ws_manager", ws)
+
+    result = await inv._llm_call(
+        system="sys",
+        messages=[{"role": "user", "content": "hi"}],
+        tools=[{"name": "noop"}],
+        turn_id="turn-abc",
+    )
+
+    # Final message returned verbatim (signed-thinking-block preservation).
+    assert result is final
+
+    # Extended thinking enabled with the documented budget.
+    assert stream.kwargs["thinking"] == {
+        "type": "enabled",
+        "budget_tokens": inv._THINKING_BUDGET,
+    }
+    # max_tokens leaves room above the thinking budget (Anthropic requires
+    # max_tokens > thinking.budget_tokens).
+    assert stream.kwargs["max_tokens"] > inv._THINKING_BUDGET
+
+    # Exactly two thinking_delta frames broadcast — text_delta does NOT fan out.
+    deltas = [(t, p) for t, p in ws.events if t == "thinking_delta"]
+    assert len(deltas) == 2
+
+    # Frame shape matches EventBusMap.thinking_delta in ws.types.ts.
+    for _, payload in deltas:
+        assert set(payload.keys()) == {"agent", "content", "turn_id"}
+        assert payload["agent"] == "investigator"
+        assert payload["turn_id"] == "turn-abc"
+    assert deltas[0][1]["content"] == "The vibration "
+    assert deltas[1][1]["content"] == "peak suggests bearing wear."
+
+
+@pytest.mark.asyncio
+async def test_llm_call_skips_empty_thinking_chunk(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Empty/None thinking chunks must not produce noisy frames."""
+    final = _FakeMessage(content=[], stop_reason="end_turn")
+    events = [
+        _FakeStreamEvent(
+            type="content_block_delta",
+            delta=_FakeStreamDelta(type="thinking_delta", thinking=""),
+        ),
+        _FakeStreamEvent(
+            type="content_block_delta",
+            delta=_FakeStreamDelta(type="thinking_delta", thinking=None),
+        ),
+    ]
+    stream = _FakeMessagesStream(events=events, final_message=final)
+
+    class _AntStub:
+        class messages:
+            pass
+
+    _AntStub.messages.stream = stream  # type: ignore[attr-defined]
+
+    ws = _FakeWS()
+    monkeypatch.setattr(inv, "anthropic", _AntStub)
+    monkeypatch.setattr(inv, "ws_manager", ws)
+
+    await inv._llm_call(system="s", messages=[], tools=[], turn_id="t")
+
+    assert [t for t, _ in ws.events if t == "thinking_delta"] == []