feat(fastapi): cap step payloads and optimize token streaming

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'")

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