Include input and output messages in weave observability traces (#1050)

This PR ensures the input message and output message are displayed in the root observability trace in Weave when workflows are executed using either `nat run` or `nat serve` with NAT-UI. 

## With `nat run`:
<img width="1022" height="112" alt="image" src="https://github.com/user-attachments/assets/edf12d4d-d458-4779-b54b-9c508533aa93" />


## With `nat serve`
### `/chat`
<img width="1088" height="113" alt="image" src="https://github.com/user-attachments/assets/4e59f587-4208-4949-a63b-13eb42e53f8e" />

### `/generate`
<img width="1030" height="112" alt="image" src="https://github.com/user-attachments/assets/f17117a9-a8ec-40a6-81a0-33d7d6335dde" />


###  `/chat/stream`:
For streaming endpoints, a preview of the first few tokens is collected for display in Weave

<img width="1127" height="119" alt="image" src="https://github.com/user-attachments/assets/b46d65bd-bfd8-499b-a9d7-c6f364e8f0a4" />

### `/generate/stream`
<img width="940" height="122" alt="image" src="https://github.com/user-attachments/assets/e320d67b-b93d-43b8-b2b9-ea1ab944f009" />

## Websocket Schemas
### `chat` 

<img width="1138" height="114" alt="image" src="https://github.com/user-attachments/assets/988dea72-7957-48b2-8c98-6fa9dee880c0" />

### `chat_stream`

<img width="1405" height="117" alt="image" src="https://github.com/user-attachments/assets/c2c87012-5d42-4580-9ade-d601fa3209ea" />

### `generate`

<img width="939" height="116" alt="image" src="https://github.com/user-attachments/assets/719658cc-8e5e-4013-8c8c-f17a7ae0b3ea" />

### `generate_stream`

<img width="979" height="117" alt="image" src="https://github.com/user-attachments/assets/cb0e9c84-92b7-46a1-b3c6-ee133425a589" />



Closes #1041 

## By Submitting this PR I confirm:
- I am familiar with the [Contributing Guidelines](https://github.com/NVIDIA/NeMo-Agent-Toolkit/blob/develop/docs/source/resources/contributing.md).
- We require that all contributors "sign-off" on their commits. This certifies that the contribution is your original work, or you have rights to submit it under the same license, or a compatible license.
  - Any contribution which contains commits that are not Signed-Off will not be accepted.
- When the PR is ready for review, new or existing tests cover these changes.
- When the PR is ready for review, the documentation is up to date with these changes.



## Summary by CodeRabbit

* **New Features**
  * Workflow start/end events now include structured input/output data.
  * Streaming workflows include an output preview (up to 50 items) on completion.

* **Improvements**
  * Better extraction of inputs from websocket-style messages.
  * Enhanced parsing to surface message/content details alongside raw outputs.
  * Added automatic truncation utility for consistent, readable output previews.

Authors:
  - Patrick Chin (https://github.com/thepatrickchin)

Approvers:
  - Will Killian (https://github.com/willkill07)

URL: #1050

Author: thepatrickchin

packages/nvidia_nat_weave/src/nat/plugins/weave/weave_exporter.py

@@ -16,12 +16,14 @@
 import logging
 from collections.abc import Generator
 from contextlib import contextmanager
+from typing import Any
 
 from nat.data_models.intermediate_step import IntermediateStep
 from nat.data_models.span import Span
 from nat.observability.exporter.base_exporter import IsolatedAttribute
 from nat.observability.exporter.span_exporter import SpanExporter
 from nat.utils.log_utils import LogFilter
+from nat.utils.string_utils import truncate_string
 from nat.utils.type_utils import override
 from weave.trace.context import weave_client_context
 from weave.trace.context.call_context import get_current_call
@@ -152,6 +154,7 @@ def _create_weave_call(self, step: IntermediateStep, span: Span) -> Call:
             try:
                 # Add the input to the Weave call
                 inputs["input"] = step.payload.data.input
+                self._extract_input_message(step.payload.data.input, inputs)
             except Exception:
                 # If serialization fails, use string representation
                 inputs["input"] = str(step.payload.data.input)
@@ -176,6 +179,76 @@ def _create_weave_call(self, step: IntermediateStep, span: Span) -> Call:
 
         return call
 
+    def _extract_input_message(self, input_data: Any, inputs: dict[str, Any]) -> None:
+        """
+        Extract message content from input data and add to inputs dictionary.
+        Also handles websocket mode where message is located at messages[0].content[0].text.
+
+        Args:
+            input_data: The raw input data from the request
+            inputs: Dictionary to populate with extracted message content
+        """
+        # Extract message content if input has messages attribute
+        messages = getattr(input_data, 'messages', [])
+        if messages:
+            content = messages[0].content
+            if isinstance(content, list) and content:
+                inputs["input_message"] = getattr(content[0], 'text', content[0])
+            else:
+                inputs["input_message"] = content
+
+    def _extract_output_message(self, output_data: Any, outputs: dict[str, Any]) -> None:
+        """
+        Extract message content from various response formats and add a preview to the outputs dictionary.
+        No data is added to the outputs dictionary if the output format is not supported.
+
+        Supported output formats for message content include:
+
+        - output.choices[0].message.content     /chat endpoint
+        - output.value                          /generate endpoint
+        - output[0].choices[0].message.content  chat WS schema
+        - output[0].choices[0].delta.content    chat_stream WS schema, /chat/stream endpoint
+        - output[0].value                       generate & generate_stream WS schema, /generate/stream endpoint
+
+        Args:
+            output_data: The raw output data from the response
+            outputs: Dictionary to populate with extracted message content.
+        """
+        # Handle choices-keyed output object for /chat completion endpoint
+        choices = getattr(output_data, 'choices', None)
+        if choices:
+            outputs["output_message"] = truncate_string(choices[0].message.content)
+            return
+
+        # Handle value-keyed output object for union types common for /generate completion endpoint
+        value = getattr(output_data, 'value', None)
+        if value:
+            outputs["output_message"] = truncate_string(value)
+            return
+
+        # Handle list-based outputs (streaming or websocket)
+        if not isinstance(output_data, list) or not output_data:
+            return
+
+        choices = getattr(output_data[0], 'choices', None)
+        if choices:
+            # chat websocket schema
+            message = getattr(choices[0], 'message', None)
+            if message:
+                outputs["output_message"] = truncate_string(getattr(message, 'content', None))
+                return
+
+            # chat_stream websocket schema and /chat/stream completion endpoint
+            delta = getattr(choices[0], 'delta', None)
+            if delta:
+                outputs["output_preview"] = truncate_string(getattr(delta, 'content', None))
+                return
+
+        # generate & generate_stream websocket schema, and /generate/stream completion endpoint
+        value = getattr(output_data[0], 'value', None)
+        if value:
+            outputs["output_preview"] = truncate_string(str(value))
+
     def _finish_weave_call(self, step: IntermediateStep) -> None:
         """
         Finish a previously created Weave call.
@@ -196,6 +269,7 @@ def _finish_weave_call(self, step: IntermediateStep) -> None:
             try:
                 # Add the output to the Weave call
                 outputs["output"] = step.payload.data.output
+                self._extract_output_message(step.payload.data.output, outputs)
             except Exception:
                 # If serialization fails, use string representation
                 outputs["output"] = str(step.payload.data.output)

src/nat/runtime/runner.py

@@ -170,7 +170,8 @@ async def result(self, to_type: type | None = None):
                     IntermediateStepPayload(UUID=workflow_step_uuid,
                                             event_type=IntermediateStepType.WORKFLOW_START,
                                             name=workflow_name,
-                                            metadata=start_metadata))
+                                            metadata=start_metadata,
+                                            data=StreamEventData(input=self._input_message)))
 
                 result = await self._entry_fn.ainvoke(self._input_message, to_type=to_type)  # type: ignore
 
@@ -249,9 +250,15 @@ async def result_stream(self, to_type: type | None = None):
                     IntermediateStepPayload(UUID=workflow_step_uuid,
                                             event_type=IntermediateStepType.WORKFLOW_START,
                                             name=workflow_name,
-                                            metadata=start_metadata))
+                                            metadata=start_metadata,
+                                            data=StreamEventData(input=self._input_message)))
+
+                # Collect preview of streaming results for the WORKFLOW_END event
+                output_preview = []
 
                 async for m in self._entry_fn.astream(self._input_message, to_type=to_type):  # type: ignore
+                    if len(output_preview) < 50:
+                        output_preview.append(m)
                     yield m
 
                 # Emit WORKFLOW_END
@@ -265,7 +272,8 @@ async def result_stream(self, to_type: type | None = None):
                     IntermediateStepPayload(UUID=workflow_step_uuid,
                                             event_type=IntermediateStepType.WORKFLOW_END,
                                             name=workflow_name,
-                                            metadata=end_metadata))
+                                            metadata=end_metadata,
+                                            data=StreamEventData(output=output_preview)))
                 self._state = RunnerState.COMPLETED
 
                 # Close the intermediate stream

src/nat/utils/string_utils.py

@@ -36,3 +36,19 @@ def convert_to_str(value: Any) -> str:
         return str(value)
     else:
         raise ValueError(f"Unsupported type for conversion to string: {type(value)}")
+
+
+def truncate_string(text: str | None, max_length: int = 100) -> str | None:
+    """
+    Truncate a string to a maximum length, adding ellipsis if truncated.
+
+    Args:
+        text: The text to truncate (can be None)
+        max_length: Maximum allowed length (default: 100)
+
+    Returns:
+        The truncated text with ellipsis if needed, or None if input was None
+    """
+    if not text or len(text) <= max_length:
+        return text
+    return text[:max_length - 3] + "..."