ouyangbo/hermes-agent
1
0
Fork 0
Code Issues Pull Requests Actions 9 Packages Projects Releases Wiki Activity

fix(relay): build the chat-completions Relay response from collector-observed chunks

Browse Source 
Relay invokes its finalizer as soon as the provider stream ends, concurrently with
Hermes' consumer thread, so any finalizer that reads the consumer loop's closures
races the last chunk. #103104 moved `usage` onto the collector; the same race still
truncated the final tool-call `arguments` delta, dropped the last content delta, and
downgraded `finish_reason` to "stop" in Relay's annotated response (reproduced ~50%
of runs). Every other Relay integration (Anthropic, Bedrock, Codex) already
rebuilds from collector-observed chunks; this makes chat_completions match via
`_RelayChatAccumulator` and removes the single-field nonlocal stopgap.

Test: the ordering harness is shared, and a second invariant test pins the
tool-call/finish_reason case (fails on the stopgap, passes here). Hermes' own
returned response was never affected.
This commit is contained in:
kshitijk4poor
2026-09-05 10:25:25 +05:30
committed by kshitij
parent 2cf5e68c5d
commit 74de0fd4fe
3 changed files with 158 additions and 98 deletions
Download Patch File Download Diff File Expand all files Collapse all files
agent/chat_completion_helpers.py
+4 -19
View File
@@ -2922,31 +2922,16 @@ class _StreamingCall:
tool_calls = _ToolCallAccumulator()
tool_calls_acc = tool_calls.acc
finish_reason = model_name = usage_obj = None
relay_observed_usage = None
role = "assistant"
_diag = self._new_diag()
self._writer_token = self._attempt_request_client = self._attempt_stream_response = None
from agent.chat_completion_helpers_relay import RelayChatAccumulator
relay_response = RelayChatAccumulator()
def _open_stream(next_api_kwargs: dict[str, Any]):
timeout = _httpx.Timeout(connect=conn_cap, read=read_timeout, write=base_timeout, pool=conn_cap)
return self._open_chat_stream({**next_api_kwargs, "stream": True, "timeout": timeout})
def _relay_final_response() -> dict[str, Any]:
tool_calls.materialize()
message = {"role": role, "content": "".join(content_parts) or None,
"reasoning_content": "".join(reasoning_parts) or None,
"tool_calls": [tool_calls_acc[i] for i in sorted(tool_calls_acc)] or None}
return {"model": model_name,
"usage": relay_observed_usage if relay_observed_usage is not None else usage_obj,
"choices": [{"message": message, "finish_reason": finish_reason or "stop"}]}
def _capture_relay_usage(chunk: Any) -> None:
# Relay orders this collector before its finalizer; the consumer may not
# have copied a terminal usage-only frame into ``usage_obj`` yet.
nonlocal relay_observed_usage
if isinstance(chunk, dict) and chunk.get("usage") is not None:
relay_observed_usage = chunk["usage"]
def _flush_pending_stream_text():
pending_parts = list(pending_text_parts)
pending_text_parts.clear()
@@ -2955,8 +2940,8 @@ class _StreamingCall:
from agent import relay_llm
stream = self._set_managed_stream(relay_llm.stream(self.api_kwargs, _open_stream,
**_relay_stream_identity(self.agent, "provider"), finalizer=_relay_final_response,
on_stream_created=self._chat_stream_created, on_chunk=_capture_relay_usage,
**_relay_stream_identity(self.agent, "provider"), finalizer=relay_response.finalize,
on_stream_created=self._chat_stream_created, on_chunk=relay_response.observe,
accept_chunk=lambda chunk: self._accept_chat_chunk(stream_attempt_id, chunk),
completed_response_predicate=lambda value: hasattr(value, "choices"),
metadata=_relay_stream_metadata(self.agent, "chat_completions"), defer_logical_completion=True))
agent/chat_completion_helpers_relay.py
+74
View File
@@ -0,0 +1,74 @@
"""Relay-side accumulator for the chat_completions streaming wire.
Relay invokes its collector for every post-intercept chunk and then its finalizer as soon
as the provider stream ends — concurrently with Hermes' consumer thread, which may not have
read the last chunk yet. The finalizer therefore builds Relay's recorded response from
collector-observed state only, never from the consumer loop's closures. Sibling of
``relay_llm.AnthropicStreamAccumulator``; Bedrock and Codex follow the same contract.
"""
from __future__ import annotations
from types import SimpleNamespace
from typing import Any
from agent.chat_completion_helpers import _ToolCallAccumulator
from agent.message_content import flatten_message_text
from agent.reasoning_summaries import separate_glued_reasoning_blocks
def _tool_call_delta_view(tc_delta: Any) -> Any:
"""Attribute view of a JSON tool-call delta for ``_ToolCallAccumulator.feed`` (written
against SDK objects). Only ``function`` is wrapped: ``feed`` passes ``extra_content``
(a dict) straight through ``_dump_if_model``, so a recursive view would corrupt it."""
if not isinstance(tc_delta, dict):
return tc_delta
function = tc_delta.get("function")
return SimpleNamespace(**{**tc_delta,
"function": SimpleNamespace(**function) if isinstance(function, dict) else function})
class RelayChatAccumulator:
"""Rebuild a chat.completion from Relay's post-intercept chunk dicts."""
def __init__(self) -> None:
self._content: list[str] = []
self._reasoning: list[str] = []
self._tool_calls = _ToolCallAccumulator()
self._model = self._usage = self._finish_reason = None
self._role = "assistant"
def observe(self, chunk: Any) -> None:
if not isinstance(chunk, dict):
return
self._model = chunk.get("model") or self._model
if chunk.get("usage"):
self._usage = chunk["usage"]
choices = chunk.get("choices") or []
choice = choices[0] if choices else None # Hermes never requests n>1
if not isinstance(choice, dict):
return
self._finish_reason = choice.get("finish_reason") or self._finish_reason
delta = choice.get("delta")
if not isinstance(delta, dict):
return
if delta.get("role"):
self._role = delta["role"]
text = flatten_message_text(delta.get("content"), sep="")
if text:
self._content.append(text)
reasoning = delta.get("reasoning_content") or delta.get("reasoning")
if reasoning:
self._reasoning.append(separate_glued_reasoning_blocks(
self._reasoning[-1] if self._reasoning else "", reasoning))
for tc_delta in delta.get("tool_calls") or []:
self._tool_calls.feed(_tool_call_delta_view(tc_delta))
def finalize(self) -> dict[str, Any]:
acc = self._tool_calls.materialize()
message = {"role": self._role, "content": "".join(self._content) or None,
"reasoning_content": "".join(self._reasoning) or None,
"tool_calls": [acc[i] for i in sorted(acc)] or None}
# "stop" also covers Nous Portal ``lastOne`` usage frames, which carry no finish_reason.
return {"model": self._model, "usage": self._usage,
"choices": [{"message": message, "finish_reason": self._finish_reason or "stop"}]}
tests/e2e/test_relay_native_openai_stream.py
+80 -79
View File
@@ -1,4 +1,10 @@
"""Native OpenAI SDK streaming through Relay's managed execution path."""
"""Native OpenAI SDK streaming through Relay's managed execution path.
Relay runs its finalizer as soon as the provider stream ends — concurrently with Hermes'
consumer thread, which may not have processed the last chunk yet. Each test forces that
ordering deterministically (finalizer runs BEFORE the consumer sees a chosen chunk) and
asserts Relay's LLM end event still records the full response.
"""
from __future__ import annotations
@@ -6,9 +12,17 @@ import threading
import pytest
_CHUNK_PREFIX = b'data: {"id":"chatcmpl-test","object":"chat.completion.chunk","created":1,"model":"test/model",'
def test_openai_stream_usage_reaches_relay_parent_event(tmp_path, monkeypatch):
"""A trailing usage-only chunk is retained on Relay's parent LLM event."""
def _sse(*chunk_bodies: bytes) -> bytes:
return b"".join(_CHUNK_PREFIX + body + b"}\n\n" for body in chunk_bodies) + b"data: [DONE]\n\n"
def _stream_through_relay(tmp_path, monkeypatch, response_body: bytes, *, finalize_before):
"""Stream ``response_body`` through Relay; Relay's finalizer is forced to complete before
the consumer thread processes the first chunk matching ``finalize_before(chunk)``.
Returns ``(hermes_result, relay_llm_end_event)``."""
httpx = pytest.importorskip("httpx")
nemo_relay = pytest.importorskip("nemo_relay")
openai = pytest.importorskip("openai")
@@ -18,55 +32,26 @@ def test_openai_stream_usage_reaches_relay_parent_event(tmp_path, monkeypatch):
monkeypatch.setenv("HERMES_HOME", str(tmp_path / "hermes-home"))
monkeypatch.setenv("HERMES_STREAM_RETRIES", "0")
response_body = b"""data: {"id":"chatcmpl-test","object":"chat.completion.chunk","created":1,"model":"test/model","choices":[{"index":0,"delta":{"role":"assistant","content":"done"},"finish_reason":null}]}
data: {"id":"chatcmpl-test","object":"chat.completion.chunk","created":1,"model":"test/model","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}
data: {"id":"chatcmpl-test","object":"chat.completion.chunk","created":1,"model":"test/model","choices":[],"usage":{"prompt_tokens":100,"completion_tokens":10,"total_tokens":110}}
data: [DONE]
"""
def respond(request):
return httpx.Response(
200,
headers={"content-type": "text/event-stream"},
content=response_body,
request=request,
)
return httpx.Response(200, headers={"content-type": "text/event-stream"},
content=response_body, request=request)
client = openai.OpenAI(
api_key="test-key",
base_url="https://example.com/v1",
http_client=httpx.Client(transport=httpx.MockTransport(respond)),
)
client = openai.OpenAI(api_key="test-key", base_url="https://example.com/v1",
http_client=httpx.Client(transport=httpx.MockTransport(respond)))
relay_runtime._reset_for_tests()
agent = AIAgent(
api_key="test-key",
base_url="https://example.com/v1",
provider="test-provider",
model="test/model",
quiet_mode=True,
skip_context_files=True,
skip_memory=True,
)
agent = AIAgent(api_key="test-key", base_url="https://example.com/v1", provider="test-provider",
model="test/model", quiet_mode=True, skip_context_files=True, skip_memory=True)
agent.api_mode = "chat_completions"
agent.session_id = "openai-relay-session"
agent._interrupt_requested = False
agent._create_request_openai_client = lambda *args, **kwargs: client
lease = relay_runtime.SESSION_COORDINATOR.acquire_conversation(
profile_key=relay_runtime.current_profile_key(),
session_id=agent.session_id,
platform="cli",
)
profile_key=relay_runtime.current_profile_key(), session_id=agent.session_id, platform="cli")
turn = relay_runtime.SESSION_COORDINATOR.begin_turn(
lease,
turn_id="openai-relay-turn",
task_id="openai-relay-task",
)
lease, turn_id="openai-relay-turn", task_id="openai-relay-task")
consumer = "test.openai_relay"
subscriber_name = "test.openai_stream_usage"
subscriber_name = "test.openai_stream"
events = []
relay_finalizer_started = threading.Event()
allow_relay_finalizer = threading.Event()
@@ -75,43 +60,30 @@ data: [DONE]
def run_synchronized_relay_finalizer(managed_stream, attempt):
relay_finalizer_started.set()
assert allow_relay_finalizer.wait(5), (
"consumer did not release Relay's finalizer"
)
assert allow_relay_finalizer.wait(5), "consumer did not release Relay's finalizer"
try:
return run_relay_finalizer(managed_stream, attempt)
finally:
relay_finalizer_finished.set()
monkeypatch.setattr(
relay_llm.ManagedLlmStream,
"_relay_finalizer",
run_synchronized_relay_finalizer,
)
monkeypatch.setattr(relay_llm.ManagedLlmStream, "_relay_finalizer", run_synchronized_relay_finalizer)
parse_choiceless_chunk = chat_completion_helpers._StreamingCall._choiceless_chunk
count_chunk = chat_completion_helpers._StreamingCall._count_chunk
def parse_usage_after_relay_finalizes(chunk, finish_reason):
if not chunk.choices and getattr(chunk, "usage", None) is not None:
# Force Relay to finalize before the consumer copies the usage frame.
def count_chunk_after_relay_finalizes(self, diag, chunk):
# ``_count_chunk`` is the first thing the consumer does with every chunk.
if finalize_before(chunk):
assert relay_finalizer_started.wait(5), "Relay's finalizer did not start"
allow_relay_finalizer.set()
assert relay_finalizer_finished.wait(5), "Relay's finalizer did not finish"
return parse_choiceless_chunk(chunk, finish_reason)
return count_chunk(self, diag, chunk)
monkeypatch.setattr(
chat_completion_helpers._StreamingCall,
"_choiceless_chunk",
staticmethod(parse_usage_after_relay_finalizes),
)
monkeypatch.setattr(chat_completion_helpers._StreamingCall, "_count_chunk", count_chunk_after_relay_finalizes)
lease.host.retain_managed_execution(consumer)
lease.host.relay.subscribers.register(subscriber_name, events.append)
try:
result = agent._interruptible_streaming_api_call({
"model": "test/model",
"messages": [{"role": "user", "content": "hi"}],
})
"model": "test/model", "messages": [{"role": "user", "content": "hi"}]})
lease.host.relay.subscribers.flush()
finally:
lease.host.relay.subscribers.deregister(subscriber_name)
@@ -121,23 +93,52 @@ data: [DONE]
relay_runtime._reset_for_tests()
client.close()
assert result.usage is not None
assert result.usage.prompt_tokens == 100
assert result.usage.completion_tokens == 10
assert result.usage.total_tokens == 110
llm_end_events = [
event
for event in events
if isinstance(event, nemo_relay.ScopeEvent)
and event.name == "openai.chat_completions"
and event.category == "llm"
and event.scope_category == "end"
event for event in events
if isinstance(event, nemo_relay.ScopeEvent) and event.name == "openai.chat_completions"
and event.category == "llm" and event.scope_category == "end"
]
assert len(llm_end_events) == 1
llm_end = llm_end_events[0]
assert llm_end.annotated_response is not None
assert llm_end_events[0].annotated_response is not None
return result, llm_end_events[0]
def test_openai_stream_usage_reaches_relay_parent_event(tmp_path, monkeypatch):
"""A trailing usage-only chunk is retained on Relay's parent LLM event."""
body = _sse(
b'"choices":[{"index":0,"delta":{"role":"assistant","content":"done"},"finish_reason":null}]',
b'"choices":[{"index":0,"delta":{},"finish_reason":"stop"}]',
b'"choices":[],"usage":{"prompt_tokens":100,"completion_tokens":10,"total_tokens":110}',
)
result, llm_end = _stream_through_relay(
tmp_path, monkeypatch, body,
finalize_before=lambda chunk: not chunk.choices and getattr(chunk, "usage", None) is not None)
assert result.usage is not None
assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (100, 10, 110)
assert llm_end.annotated_response.usage == {
"prompt_tokens": 100,
"completion_tokens": 10,
"total_tokens": 110,
}
"prompt_tokens": 100, "completion_tokens": 10, "total_tokens": 110}
assert llm_end.annotated_response.message == "done"
def test_openai_stream_final_tool_call_delta_reaches_relay_parent_event(tmp_path, monkeypatch):
"""The last chunk's tool-call arguments and finish_reason are retained on Relay's parent
LLM event — the same finalizer-before-consumer race as the usage frame, without one."""
body = _sse(
b'"choices":[{"index":0,"delta":{"role":"assistant","content":null,"tool_calls":[{"index":0,'
b'"id":"call_1","type":"function","function":{"name":"read_file","arguments":"{\\"path\\": "}}]},'
b'"finish_reason":null}]',
b'"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"\\"/tmp/x\\"}"}}]},'
b'"finish_reason":"tool_calls"}]',
)
result, llm_end = _stream_through_relay(
tmp_path, monkeypatch, body,
finalize_before=lambda chunk: bool(chunk.choices) and chunk.choices[0].finish_reason == "tool_calls")
hermes_call = result.choices[0].message.tool_calls[0]
assert (hermes_call.function.name, hermes_call.function.arguments) == ("read_file", '{"path": "/tmp/x"}')
assert result.choices[0].finish_reason == "tool_calls"
assert llm_end.annotated_response.message is None
(relay_call,) = llm_end.annotated_response.tool_calls
assert (relay_call["name"], relay_call["arguments"]) == ("read_file", {"path": "/tmp/x"})
assert llm_end.annotated_response.finish_reason == "tool_use"