diff --git a/packages/nvidia_nat_core/src/nat/data_models/step_adaptor.py b/packages/nvidia_nat_core/src/nat/data_models/step_adaptor.py index 9430a269f7..3d438206a3 100644 --- a/packages/nvidia_nat_core/src/nat/data_models/step_adaptor.py +++ b/packages/nvidia_nat_core/src/nat/data_models/step_adaptor.py @@ -33,24 +33,49 @@ class StepAdaptorMode(StrEnum): class StepAdaptorConfig(BaseModel): """ - Configures how intermediate steps are filtered and normalized by the StepAdaptor. + Configures how intermediate steps are filtered and normalized by the ``StepAdaptor``. Args: - mode (StepAdaptorMode): One of: - - 'current' => pass only LLM (all LLM_* events) + TOOL_END - - 'end_events_only' => pass only LLM_END and TOOL_END - - 'custom' => pass only the events in custom_event_types + mode (StepAdaptorMode): Mode determining which events are emitted (``StepAdaptorMode.DEFAULT``, + ``StepAdaptorMode.CUSTOM``, or ``StepAdaptorMode.OFF``). custom_event_types (list[IntermediateStepType]): - If mode == 'custom', we only pass events whose event_type is in this list. + If ``mode`` is ``StepAdaptorMode.CUSTOM``, only events whose ``event_type`` is in this list are passed. Otherwise, this field is ignored. + stream_llm_tokens (bool): Whether to emit intermediate LLM token events + (``IntermediateStepType.LLM_NEW_TOKEN``). When ``False``, only ``IntermediateStepType.LLM_START`` + and ``IntermediateStepType.LLM_END`` are emitted. + max_input_length (int): Maximum character length for input fields in intermediate step payloads. + Exceeding text will be truncated. Must be greater than or equal to 0. + max_output_length (int): Maximum character length for output fields in intermediate step payloads. + Exceeding text will be truncated. Must be greater than or equal to 0. """ mode: StepAdaptorMode = StepAdaptorMode.DEFAULT custom_event_types: list[IntermediateStepType] = Field(default_factory=list) + stream_llm_tokens: bool = Field( + default=False, + description=("Whether to emit intermediate LLM token events (LLM_NEW_TOKEN). " + "When False, only LLM_START and LLM_END are emitted."), + ) + max_input_length: int = Field( + default=4000, + ge=0, + description=("Maximum character length for input fields in intermediate step payloads. " + "Exceeding text will be truncated."), + ) + max_output_length: int = Field( + default=4000, + ge=0, + description=("Maximum character length for output fields in intermediate step payloads. " + "Exceeding text will be truncated."), + ) @model_validator(mode="after") def check_custom_event_types(self) -> "StepAdaptorConfig": """ - Validates custom configurations + Validates ``StepAdaptorConfig`` when ``mode`` is ``StepAdaptorMode.CUSTOM``. + + Returns: + StepAdaptorConfig: The validated configuration instance. """ if self.mode != StepAdaptorMode.CUSTOM and self.custom_event_types: logger.warning("Ignoring custom_event_types because mode is not 'custom'") diff --git a/packages/nvidia_nat_core/src/nat/front_ends/fastapi/step_adaptor.py b/packages/nvidia_nat_core/src/nat/front_ends/fastapi/step_adaptor.py index f1325ca93a..91db949836 100644 --- a/packages/nvidia_nat_core/src/nat/front_ends/fastapi/step_adaptor.py +++ b/packages/nvidia_nat_core/src/nat/front_ends/fastapi/step_adaptor.py @@ -15,7 +15,7 @@ import html import logging -from functools import reduce +from collections import defaultdict from textwrap import dedent from nat.data_models.api_server import ResponseIntermediateStep @@ -35,18 +35,57 @@ class StepAdaptor: def __init__(self, config: StepAdaptorConfig): + """ + Initializes the ``StepAdaptor`` with configuration settings. + Args: + config (StepAdaptorConfig): The configuration governing event filtering and payload truncation. + """ self._history: list[IntermediateStep] = [] + self._llm_chunks: dict[str, list[str]] = defaultdict(list) self.config = config - def _step_matches_filter(self, step: IntermediateStep, config: StepAdaptorConfig) -> bool: + def _truncate_text(self, text: str | None, max_len: int) -> str: """ - Returns True if this intermediate step should be included (based on the config.mode). + Truncates text if it exceeds ``max_len``, appending a truncation notice. + + Args: + text (str | None): The text to truncate. + max_len (int): The maximum character length allowed. + + Returns: + str: The truncated text or empty string if input is None or empty. """ + if not text: + return "" + if max_len <= 0: + return "" + if len(text) <= max_len: + return text + + notice = f"\n... [truncated {len(text) - max_len} characters]" + if max_len > len(notice): + slice_len = max_len - len(notice) + return f"{text[:slice_len]}{notice}" + return text[:max_len] + + def _step_matches_filter(self, step: IntermediateStep, config: StepAdaptorConfig) -> bool: + """ + Determines if an intermediate step should be included based on ``config.mode``. + + Args: + step (IntermediateStep): The intermediate step to evaluate. + config (StepAdaptorConfig): The current adaptor configuration. + Returns: + bool: ``True`` if the step should be processed, ``False`` otherwise. + """ if config.mode == StepAdaptorMode.OFF: return False + if step.event_type == IntermediateStepType.LLM_NEW_TOKEN and not config.stream_llm_tokens: + return False + if config.mode == StepAdaptorMode.DEFAULT: # default existing behavior: show LLM events + TOOL_END + FUNCTION events if step.event_category == IntermediateStepCategory.LLM: @@ -64,6 +103,16 @@ def _step_matches_filter(self, step: IntermediateStep, config: StepAdaptorConfig return False def _handle_llm(self, step: IntermediateStepPayload, ancestry: InvocationNode) -> ResponseSerializable | None: + """ + Handles ``LLM_START``, ``LLM_NEW_TOKEN``, and ``LLM_END`` events. + + Args: + step (IntermediateStepPayload): The intermediate step payload. + ancestry (InvocationNode): The invocation node representing the ancestry hierarchy. + + Returns: + ResponseSerializable | None: The formatted intermediate step response, or ``None`` if skipped. + """ input_str: str | None = None output_str: str | None = None @@ -78,20 +127,23 @@ def _handle_llm(self, step: IntermediateStepPayload, ancestry: InvocationNode) - input_str = str(start_step.data.input) if step.event_type == IntermediateStepType.LLM_NEW_TOKEN: - - # Find all of the previous LLM chunks and concatenate them - output_str = reduce( - lambda x, y: x + y, - (str(x.data.chunk) - for x in self._history if x.event_type == IntermediateStepType.LLM_NEW_TOKEN and x.UUID == step.UUID), - "") + output_str = "".join(self._llm_chunks[step.UUID]) elif step.event_type == IntermediateStepType.LLM_END: - output_str = str(step.data.output) + if hasattr(step.data, "output") and step.data.output is not None and str(step.data.output).strip() != "": + output_str = str(step.data.output) + else: + output_str = "".join(self._llm_chunks.get(step.UUID, [])) + self._llm_chunks.pop(step.UUID, None) if not input_str and not output_str: return None + if input_str: + input_str = self._truncate_text(input_str, self.config.max_input_length) + if output_str: + output_str = self._truncate_text(output_str, self.config.max_output_length) + escaped_input = html.escape(input_str, quote=False) # Dont use f-strings here because the payload is markdown and screws up the dedent @@ -122,7 +174,14 @@ def _handle_llm(self, step: IntermediateStepPayload, ancestry: InvocationNode) - def _handle_tool(self, step: IntermediateStepPayload, ancestry: InvocationNode) -> ResponseSerializable | None: """ - Handles both TOOL_START and TOOL_END events + Handles both ``TOOL_START`` and ``TOOL_END`` events. + + Args: + step (IntermediateStepPayload): The intermediate step payload. + ancestry (InvocationNode): The invocation node representing the ancestry hierarchy. + + Returns: + ResponseSerializable | None: The formatted intermediate step response, or ``None`` if skipped. """ input_str: str | None = None output_str: str | None = None @@ -136,9 +195,13 @@ def _handle_tool(self, step: IntermediateStepPayload, ancestry: InvocationNode) return None input_str = str(start_step.data.input) + if input_str: + input_str = self._truncate_text(input_str, self.config.max_input_length) if step.event_type == IntermediateStepType.TOOL_END: output_str = str(step.data.output) + if output_str: + output_str = self._truncate_text(output_str, self.config.max_output_length) if not input_str and not output_str: return None @@ -177,7 +240,14 @@ def _handle_tool(self, step: IntermediateStepPayload, ancestry: InvocationNode) def _handle_function(self, step: IntermediateStepPayload, ancestry: InvocationNode) -> ResponseSerializable | None: """ - Handles the FUNCTION_START and FUNCTION_END events + Handles the ``FUNCTION_START`` and ``FUNCTION_END`` events. + + Args: + step (IntermediateStepPayload): The intermediate step payload. + ancestry (InvocationNode): The invocation node representing the ancestry hierarchy. + + Returns: + ResponseSerializable | None: The formatted intermediate step response, or ``None`` if skipped. """ input_str: str | None = None output_str: str | None = None @@ -192,6 +262,7 @@ def _handle_function(self, step: IntermediateStepPayload, ancestry: InvocationNo if not input_str: return None + input_str = self._truncate_text(input_str, self.config.max_input_length) escaped_input = html.escape(input_str, quote=False) format_input_type = "json" if is_valid_json(escaped_input) else "python" @@ -225,6 +296,7 @@ def _handle_function(self, step: IntermediateStepPayload, ancestry: InvocationNo if not output_str: return None + output_str = self._truncate_text(output_str, self.config.max_output_length) escaped_output = html.escape(output_str, quote=False) format_output_type = "json" if is_valid_json(escaped_output) else "python" @@ -237,6 +309,7 @@ def _handle_function(self, step: IntermediateStepPayload, ancestry: InvocationNo input_str = str(start_step.data) if input_str: + input_str = self._truncate_text(input_str, self.config.max_input_length) escaped_input = html.escape(input_str, quote=False) format_input_type = "json" if is_valid_json(escaped_input) else "python" input_payload = dedent(""" @@ -266,7 +339,14 @@ def _handle_function(self, step: IntermediateStepPayload, ancestry: InvocationNo def _handle_custom(self, payload: IntermediateStepPayload, ancestry: InvocationNode) -> ResponseSerializable | None: """ - Handles the CUSTOM event + Handles the ``CUSTOM`` event. + + Args: + payload (IntermediateStepPayload): The intermediate step payload. + ancestry (InvocationNode): The invocation node representing the ancestry hierarchy. + + Returns: + ResponseSerializable | None: The formatted intermediate step response, or ``None`` if skipped. """ escaped_payload = html.escape(str(payload), quote=False) escaped_payload = escaped_payload.replace("\n", "") @@ -290,13 +370,35 @@ def _handle_custom(self, payload: IntermediateStepPayload, ancestry: InvocationN return event def process(self, step: IntermediateStep) -> ResponseSerializable | None: + """ + Processes an intermediate step and returns a serialized response if matched. + + Args: + step (IntermediateStep): The intermediate step event to process. + Returns: + ResponseSerializable | None: The adapted response model if matched and processed, + or ``None`` if filtered out or an error occurred. + """ # Track the chunk self._history.append(step) payload = step.payload ancestry = step.function_ancestry + if step.event_type == IntermediateStepType.LLM_NEW_TOKEN: + if step.data: + if hasattr(step.data, "chunk") and step.data.chunk is not None: + self._llm_chunks[step.UUID].append(str(step.data.chunk)) + elif hasattr(step.data, "output") and step.data.output is not None: + self._llm_chunks[step.UUID].append(str(step.data.output)) + elif hasattr(step.data, "payload") and step.data.payload is not None: + self._llm_chunks[step.UUID].append(str(step.data.payload)) + else: + self._llm_chunks[step.UUID].append(str(step.data)) + if not self._step_matches_filter(step, self.config): + if step.event_type == IntermediateStepType.LLM_END: + self._llm_chunks.pop(step.UUID, None) return None try: @@ -315,5 +417,8 @@ def process(self, step: IntermediateStep) -> ResponseSerializable | None: except Exception as e: logger.exception("Error processing intermediate step: %s", e) + finally: + if step.event_type == IntermediateStepType.LLM_END: + self._llm_chunks.pop(step.UUID, None) return None diff --git a/packages/nvidia_nat_core/tests/nat/front_ends/fastapi/test_step_adaptor.py b/packages/nvidia_nat_core/tests/nat/front_ends/fastapi/test_step_adaptor.py index f537987cb5..7c6e93842c 100644 --- a/packages/nvidia_nat_core/tests/nat/front_ends/fastapi/test_step_adaptor.py +++ b/packages/nvidia_nat_core/tests/nat/front_ends/fastapi/test_step_adaptor.py @@ -14,6 +14,7 @@ # limitations under the License. import pytest +from pydantic import ValidationError from nat.data_models.api_server import ResponseIntermediateStep from nat.data_models.intermediate_step import IntermediateStep @@ -99,18 +100,53 @@ def _make_step(event_type: IntermediateStepType, data_input=None, data_output=No # -------------------- # Tests for DEFAULT mode # -------------------- -@pytest.mark.parametrize("event_type", [(IntermediateStepType.LLM_START)]) -def test_process_llm_events_in_default(step_adaptor_default, make_intermediate_step, event_type): +@pytest.mark.parametrize( + "event_type, expected_result", + [ + (IntermediateStepType.LLM_START, True), + (IntermediateStepType.LLM_NEW_TOKEN, False), + (IntermediateStepType.LLM_END, True), + ], +) +def test_process_llm_events_in_default(step_adaptor_default, make_intermediate_step, event_type, expected_result): + """ + In DEFAULT mode with stream_llm_tokens=False (default): + - LLM_START returns a valid ResponseIntermediateStep. + - LLM_NEW_TOKEN returns None. + - LLM_END returns a valid ResponseIntermediateStep. + All steps must be appended to _history. """ - In DEFAULT mode, LLM_START, LLM_NEW_TOKEN, and LLM_END events are processed. - We expect a valid ResponseIntermediateStep for each. - """ - step = make_intermediate_step(event_type=event_type, data_input="LLM Input", data_output="LLM Output") + step_start = make_intermediate_step( + event_type=IntermediateStepType.LLM_START, + data_input="LLM Input", + UUID="test-default-llm-uuid", + ) + + if event_type == IntermediateStepType.LLM_START: + step = step_start + elif event_type == IntermediateStepType.LLM_NEW_TOKEN: + step_adaptor_default.process(step_start) + step = make_intermediate_step( + event_type=IntermediateStepType.LLM_NEW_TOKEN, + data_output="chunk", + UUID="test-default-llm-uuid", + ) + else: + step_adaptor_default.process(step_start) + step = make_intermediate_step( + event_type=IntermediateStepType.LLM_END, + data_output="LLM Output", + UUID="test-default-llm-uuid", + ) result = step_adaptor_default.process(step) - assert result is not None, f"Expected LLM event '{event_type}' to be processed in DEFAULT mode." - assert isinstance(result, ResponseIntermediateStep) + if expected_result: + assert result is not None, f"Expected LLM event '{event_type}' to be processed in DEFAULT mode." + assert isinstance(result, ResponseIntermediateStep) + else: + assert result is None, f"Expected LLM event '{event_type}' to be filtered out in DEFAULT mode." + assert step_adaptor_default._history[-1] is step, "Step must be appended to _history." @@ -525,3 +561,358 @@ def test_function_events_in_custom_mode(step_adaptor_custom, make_intermediate_s # Steps should still be added to history assert step_adaptor_custom._history[-2] is step_start assert step_adaptor_custom._history[-1] is step_end + + +def test_truncate_text_helper(default_config): + """ + Verify that the `_truncate_text` helper correctly truncates strings exceeding limits and handles empty inputs. + """ + adaptor = StepAdaptor(config=default_config) + + assert adaptor._truncate_text(None, 10) == "" + assert adaptor._truncate_text("", 10) == "" + assert adaptor._truncate_text("hello", 10) == "hello" + assert adaptor._truncate_text("hello world", 11) == "hello world" + assert adaptor._truncate_text("hello world", 0) == "" + assert adaptor._truncate_text("hello world", -1) == "" + + truncated_short = adaptor._truncate_text("hello world", 5) + assert truncated_short == "hello" + assert len(truncated_short) <= 5 + + long_text = "A" * 100 + truncated_long = adaptor._truncate_text(long_text, 50) + assert len(truncated_long) <= 50 + assert "[truncated 50 characters]" in truncated_long + + +def test_tool_truncation(make_intermediate_step): + """ + Verify that tool inputs and outputs exceeding configured character limits are truncated. + """ + config = StepAdaptorConfig(max_input_length=50, max_output_length=55) + adaptor = StepAdaptor(config=config) + + long_input = "A" * 100 + long_output = "B" * 100 + + step_start = make_intermediate_step( + event_type=IntermediateStepType.TOOL_START, + data_input=long_input, + UUID="tool-trunc-uuid", + ) + result_start = adaptor.process(step_start) + assert result_start is not None + assert "[truncated 50 characters]" in result_start.payload + + step_end = make_intermediate_step( + event_type=IntermediateStepType.TOOL_END, + data_input=long_input, + data_output=long_output, + UUID="tool-trunc-uuid", + ) + result_end = adaptor.process(step_end) + assert result_end is not None + assert "[truncated 50 characters]" in result_end.payload + assert "[truncated 45 characters]" in result_end.payload + + +def test_function_truncation(make_intermediate_step): + """ + Verify that function inputs and outputs exceeding configured character limits are truncated. + """ + config = StepAdaptorConfig(max_input_length=50, max_output_length=60) + adaptor = StepAdaptor(config=config) + + long_input = "X" * 100 + long_output = "Y" * 100 + uuid = "func-trunc-uuid" + + step_start = make_intermediate_step( + event_type=IntermediateStepType.FUNCTION_START, + data_input=long_input, + name="long_fn", + UUID=uuid, + ) + result_start = adaptor.process(step_start) + assert result_start is not None + assert "[truncated 50 characters]" in result_start.payload + + step_end = make_intermediate_step( + event_type=IntermediateStepType.FUNCTION_END, + data_output=long_output, + name="long_fn", + UUID=uuid, + ) + result_end = adaptor.process(step_end) + assert result_end is not None + assert "[truncated 50 characters]" in result_end.payload + assert "[truncated 40 characters]" in result_end.payload + + +def test_llm_token_streaming_default_disabled(step_adaptor_default, make_intermediate_step): + """ + Verify that when `stream_llm_tokens` is `False` (default), token chunks return `None` and start/end events emit. + """ + uuid = "llm-stream-test-uuid" + step_start = make_intermediate_step( + event_type=IntermediateStepType.LLM_START, + data_input="Prompt input", + UUID=uuid, + ) + step_token = make_intermediate_step( + event_type=IntermediateStepType.LLM_NEW_TOKEN, + data_output="token chunk", + UUID=uuid, + ) + step_end = make_intermediate_step( + event_type=IntermediateStepType.LLM_END, + data_output="Final response", + UUID=uuid, + ) + + result_start = step_adaptor_default.process(step_start) + assert result_start is not None + assert isinstance(result_start, ResponseIntermediateStep) + assert "Prompt input" in result_start.payload + + result_token = step_adaptor_default.process(step_token) + assert result_token is None + + result_end = step_adaptor_default.process(step_end) + assert result_end is not None + assert isinstance(result_end, ResponseIntermediateStep) + assert "Final response" in result_end.payload + + +def test_llm_token_streaming_opt_in(make_intermediate_step): + """ + Verify that when `stream_llm_tokens` is `True`, intermediate token chunks emit cumulative streaming updates. + """ + config = StepAdaptorConfig(stream_llm_tokens=True) + adaptor = StepAdaptor(config=config) + uuid = "llm-opt-in-uuid" + + step_start = make_intermediate_step( + event_type=IntermediateStepType.LLM_START, + data_input="Prompt input", + UUID=uuid, + ) + step_token1 = make_intermediate_step( + event_type=IntermediateStepType.LLM_NEW_TOKEN, + data_output="hello ", + UUID=uuid, + ) + step_token2 = make_intermediate_step( + event_type=IntermediateStepType.LLM_NEW_TOKEN, + data_output="world", + UUID=uuid, + ) + step_end = make_intermediate_step( + event_type=IntermediateStepType.LLM_END, + data_output="hello world!", + UUID=uuid, + ) + + adaptor.process(step_start) + + result_t1 = adaptor.process(step_token1) + assert result_t1 is not None + assert "hello " in result_t1.payload + + result_t2 = adaptor.process(step_token2) + assert result_t2 is not None + assert "hello world" in result_t2.payload + + result_end = adaptor.process(step_end) + assert result_end is not None + assert "hello world!" in result_end.payload + assert uuid not in adaptor._llm_chunks + + +def test_llm_truncation(make_intermediate_step): + """ + Verify that LLM inputs and outputs exceeding configured character limits are truncated. + """ + config = StepAdaptorConfig(stream_llm_tokens=True, max_input_length=50, max_output_length=55) + adaptor = StepAdaptor(config=config) + uuid = "llm-trunc-uuid" + + long_prompt = "P" * 100 + long_output = "R" * 100 + + step_start = make_intermediate_step( + event_type=IntermediateStepType.LLM_START, + data_input=long_prompt, + UUID=uuid, + ) + result_start = adaptor.process(step_start) + assert result_start is not None + assert "[truncated 50 characters]" in result_start.payload + + step_end = make_intermediate_step( + event_type=IntermediateStepType.LLM_END, + data_output=long_output, + UUID=uuid, + ) + result_end = adaptor.process(step_end) + assert result_end is not None + assert "[truncated 50 characters]" in result_end.payload + assert "[truncated 45 characters]" in result_end.payload + + +def test_config_negative_length_validation(): + """ + Verify that `StepAdaptorConfig` raises a `ValidationError` when negative lengths are supplied. + """ + with pytest.raises(ValidationError): + StepAdaptorConfig(max_input_length=-1) + + with pytest.raises(ValidationError): + StepAdaptorConfig(max_output_length=-1) + + +def test_llm_chunk_payload_fallback(make_intermediate_step): + """ + Verify that `_handle_llm` extracts token chunks from `step.data.payload` if `chunk` and `output` are `None`. + """ + config = StepAdaptorConfig(stream_llm_tokens=True) + adaptor = StepAdaptor(config=config) + uuid = "llm-payload-fallback-uuid" + + step_start = make_intermediate_step( + event_type=IntermediateStepType.LLM_START, + data_input="Prompt input", + UUID=uuid, + ) + adaptor.process(step_start) + + step_token = make_intermediate_step( + event_type=IntermediateStepType.LLM_NEW_TOKEN, + UUID=uuid, + ) + step_token.data.payload = "chunk from payload" + + result_token = adaptor.process(step_token) + assert result_token is not None + assert "chunk from payload" in result_token.payload + + +def test_llm_chunks_cleanup_when_end_filtered(make_intermediate_step): + """ + Verify that `_llm_chunks` buffer is cleaned up upon `LLM_END` even if the end event is filtered out. + """ + config = StepAdaptorConfig( + mode=StepAdaptorMode.CUSTOM, + custom_event_types=[IntermediateStepType.LLM_START, IntermediateStepType.LLM_NEW_TOKEN], + stream_llm_tokens=True, + ) + adaptor = StepAdaptor(config=config) + uuid = "llm-filtered-end-uuid" + + step_start = make_intermediate_step( + event_type=IntermediateStepType.LLM_START, + data_input="Prompt", + UUID=uuid, + ) + step_token = make_intermediate_step( + event_type=IntermediateStepType.LLM_NEW_TOKEN, + data_output="streamed chunk", + UUID=uuid, + ) + step_end = make_intermediate_step( + event_type=IntermediateStepType.LLM_END, + data_output="Final Output", + UUID=uuid, + ) + + adaptor.process(step_start) + adaptor.process(step_token) + assert uuid in adaptor._llm_chunks + assert adaptor._llm_chunks[uuid] == ["streamed chunk"] + + result_end = adaptor.process(step_end) + assert result_end is None, "LLM_END should be filtered out in this custom mode" + assert uuid not in adaptor._llm_chunks, "_llm_chunks must be cleaned up on LLM_END even when filtered" + + +def test_llm_end_output_fallback_to_chunks(make_intermediate_step): + """ + Verify that `LLM_END` falls back to joining accumulated chunks when `step.data.output` is `None` or empty. + """ + config = StepAdaptorConfig(stream_llm_tokens=True) + adaptor = StepAdaptor(config=config) + uuid = "llm-end-fallback-uuid" + + step_start = make_intermediate_step( + event_type=IntermediateStepType.LLM_START, + data_input="Prompt input", + UUID=uuid, + ) + step_token1 = make_intermediate_step( + event_type=IntermediateStepType.LLM_NEW_TOKEN, + data_output="streamed ", + UUID=uuid, + ) + step_token2 = make_intermediate_step( + event_type=IntermediateStepType.LLM_NEW_TOKEN, + data_output="content", + UUID=uuid, + ) + step_end = make_intermediate_step( + event_type=IntermediateStepType.LLM_END, + data_output=None, + UUID=uuid, + ) + + adaptor.process(step_start) + adaptor.process(step_token1) + adaptor.process(step_token2) + assert uuid in adaptor._llm_chunks + + result_end = adaptor.process(step_end) + assert result_end is not None + assert "streamed content" in result_end.payload + assert uuid not in adaptor._llm_chunks + + +def test_llm_token_accumulation_in_default_mode(step_adaptor_default, make_intermediate_step): + """ + Verify that in default mode (stream_llm_tokens=False), tokens accumulate and emit full output on LLM_END. + """ + uuid = "llm-accum-default-uuid" + + step_start = make_intermediate_step( + event_type=IntermediateStepType.LLM_START, + data_input="Prompt input", + UUID=uuid, + ) + step_token1 = make_intermediate_step( + event_type=IntermediateStepType.LLM_NEW_TOKEN, + data_output="Hello ", + UUID=uuid, + ) + step_token2 = make_intermediate_step( + event_type=IntermediateStepType.LLM_NEW_TOKEN, + data_output="World!", + UUID=uuid, + ) + step_end = make_intermediate_step( + event_type=IntermediateStepType.LLM_END, + data_output=None, + UUID=uuid, + ) + + result_start = step_adaptor_default.process(step_start) + assert result_start is not None + + result_t1 = step_adaptor_default.process(step_token1) + assert result_t1 is None + + result_t2 = step_adaptor_default.process(step_token2) + assert result_t2 is None + + result_end = step_adaptor_default.process(step_end) + assert result_end is not None + assert "Hello World!" in result_end.payload + assert uuid not in step_adaptor_default._llm_chunks