feat: HomeFormatter parses inline entity tags instead of tool_end blocks

The supervisor LLM now embeds <type>[id1,id2]</type> entity tags in its
response text. The HomeFormatter buffers streamed tokens, detects complete
tags across chunk boundaries, and emits WsStreamBlock with entity type +
specific IDs. This replaces the old approach of emitting blocks for every
tool_end event, which dumped ALL entities regardless of relevance.

Also fixes:
- NoneType guard on metadata in _run_graph_stream (metadata can be None)
- Updated _HOME_SYSTEM prompt with entity tag instructions
- Updated all affected tests
This commit is contained in:
2026-03-12 00:01:06 +01:00
parent 92716cb89a
commit 617a17db40
4 changed files with 190 additions and 51 deletions

View File

@@ -6,9 +6,10 @@ Consumes ``(event_type, data)`` tuples yielded by ``deep_agent.run_*_stream()``:
* ``("mutations", list)`` — collected CRUD mutations for ``stream_end``
HomeFormatter:
* Sniffs ``tool_end`` events → emits ``WsStreamBlock`` (entity_ref with raw data)
* Streams text tokens → emits ``WsStreamText``
* Attaches mutations → injects into ``WsStreamEnd``
* Buffers text tokens and parses inline entity tags
``<type>[id1,id2]</type>`` → emits ``WsStreamBlock`` (entity_ref with IDs)
* Streams surrounding text → emits ``WsStreamText``
* Attaches mutations → injects into ``WsStreamEnd``
FloatingFormatter:
* Sniffs first ``tool_end`` name → emits ``WsFloatingDomain``
@@ -19,6 +20,7 @@ FloatingFormatter:
from __future__ import annotations
import logging
import re
from collections.abc import AsyncGenerator
from typing import Any
@@ -40,20 +42,91 @@ _AGENT_DOMAIN: dict[str, str] = {
"project_agent": "projects",
}
# Regex for complete inline entity tags: <task>[id1,id2]</task>
_ENTITY_TAG_RE = re.compile(
r"<(task|project|note|timeline)>\[([^\]]+)\]</\1>"
)
# Tag name → plural entity type for the WsStreamBlock data
_TAG_ENTITY: dict[str, str] = {
"task": "tasks",
"project": "projects",
"note": "notes",
"timeline": "timelines",
}
WsFrame = WsStreamStart | WsStreamText | WsStreamBlock | WsStreamEnd | WsFloatingDomain
class HomeFormatter:
"""Consumes a deep-agent event stream and yields WS frames for the Home view.
``tool_end`` events from sub-agents are emitted as ``WsStreamBlock``
(entity_ref) so the client can render structured data. Text tokens are
forwarded as ``WsStreamText``. Mutations are attached to ``WsStreamEnd``.
The supervisor's response contains inline entity tags like
``<project>[abc-123]</project>``. This formatter detects them,
emits ``WsStreamBlock(block_type="entity_ref")`` with the entity
type and IDs, and forwards surrounding text as ``WsStreamText``.
Mutations are attached to ``WsStreamEnd``.
"""
def __init__(self, request_id: str) -> None:
self.request_id = request_id
self._mutations: list[dict] = []
self._buffer: str = ""
def _flush_buffer(self, force: bool = False):
"""Extract complete entity tags and text from the buffer.
Yields (frame_type, data) pairs:
("text", str) — plain text to send as WsStreamText
("block", dict) — entity_ref block to send as WsStreamBlock
When *force* is True (end of stream), the entire buffer is flushed.
Otherwise, text after the last unmatched ``<`` is held back in case
it is the start of an entity tag arriving across token boundaries.
"""
buf = self._buffer
while True:
m = _ENTITY_TAG_RE.search(buf)
if not m:
break
# Text before the tag
before = buf[: m.start()]
if before:
yield ("text", before)
# The entity tag itself → a block
tag_type = m.group(1)
raw_ids = m.group(2)
ids = [i.strip() for i in raw_ids.split(",") if i.strip()]
yield (
"block",
{
"entity": _TAG_ENTITY[tag_type],
"ids": ids,
},
)
buf = buf[m.end() :]
if force:
# End of stream — flush everything that remains
if buf:
yield ("text", buf)
self._buffer = ""
else:
# Keep a potential partial tag (text after last '<') in the buffer
last_lt = buf.rfind("<")
if last_lt != -1:
safe = buf[:last_lt]
if safe:
yield ("text", safe)
self._buffer = buf[last_lt:]
else:
if buf:
yield ("text", buf)
self._buffer = ""
async def format(
self,
@@ -64,22 +137,36 @@ class HomeFormatter:
async for event_type, data in event_stream:
if event_type == "token":
if data:
yield WsStreamText(request_id=self.request_id, chunk=data)
elif event_type == "tool_end":
# Sub-agent finished — emit its result as an entity_ref block
name = data.get("name", "")
entity = _AGENT_DOMAIN.get(name)
if entity:
yield WsStreamBlock(
request_id=self.request_id,
block_type="entity_ref",
data={"entity": entity, "result": data.get("result", "")},
)
self._buffer += data
for ftype, fdata in self._flush_buffer():
if ftype == "text":
yield WsStreamText(
request_id=self.request_id, chunk=fdata
)
elif ftype == "block":
yield WsStreamBlock(
request_id=self.request_id,
block_type="entity_ref",
data=fdata,
)
elif event_type == "mutations":
self._mutations = data or []
# tool_end events are intentionally ignored — the supervisor
# embeds relevant entity IDs inline via <type>[ids]</type> tags.
# Flush any remaining buffer content
for ftype, fdata in self._flush_buffer(force=True):
if ftype == "text":
yield WsStreamText(request_id=self.request_id, chunk=fdata)
elif ftype == "block":
yield WsStreamBlock(
request_id=self.request_id,
block_type="entity_ref",
data=fdata,
)
yield WsStreamEnd(
request_id=self.request_id,
mutations=[