diff --git a/agent/chat_completion_helpers.py b/agent/chat_completion_helpers.py index 70b002142e..366931ed95 100644 --- a/agent/chat_completion_helpers.py +++ b/agent/chat_completion_helpers.py @@ -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)) diff --git a/agent/chat_completion_helpers_relay.py b/agent/chat_completion_helpers_relay.py new file mode 100644 index 0000000000..37dd8b850f --- /dev/null +++ b/agent/chat_completion_helpers_relay.py @@ -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"}]} diff --git a/tests/e2e/test_relay_native_openai_stream.py b/tests/e2e/test_relay_native_openai_stream.py index 4f505c0a67..0ee25472e2 100644 --- a/tests/e2e/test_relay_native_openai_stream.py +++ b/tests/e2e/test_relay_native_openai_stream.py @@ -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"