feat(core): add thought_description for simplified progress display

packages/nvidia_nat_core/src/nat/builder/function.py

@@ -76,6 +76,7 @@ def __init__(self,
         self.description = description
         self.instance_name = instance_name or config.type
         self.display_name = config.name or self.instance_name
+        self.thought_description = config.thought_description
         self._context = Context.get()
         self._configured_middleware: tuple[Middleware, ...] = tuple()
         self._middlewared_single: _InvokeFnT | None = None
@@ -192,8 +193,9 @@ async def ainvoke(self, value: InputT | typing.Any, to_type: type | None = None)
             If the output of the function cannot be converted to the specified type.
         """
 
-        with self._context.push_active_function(self.instance_name,
-                                                input_data=value) as manager:  # Set the current invocation context
+        metadata = {"thought_description": self.thought_description} if self.thought_description else None
+        with self._context.push_active_function(self.instance_name, input_data=value,
+                                                metadata=metadata) as manager:  # Set the current invocation context
             try:
                 converted_input: InputT = self._convert_input(value)
 
@@ -285,7 +287,8 @@ async def astream(self, value: InputT | typing.Any, to_type: type | None = None)
             If the output of the function cannot be converted to the specified type (when `to_type` is specified).
         """
 
-        with self._context.push_active_function(self.instance_name, input_data=value) as manager:
+        metadata = {"thought_description": self.thought_description} if self.thought_description else None
+        with self._context.push_active_function(self.instance_name, input_data=value, metadata=metadata) as manager:
             try:
                 converted_input: InputT = self._convert_input(value)
 

packages/nvidia_nat_core/src/nat/data_models/api_server.py

@@ -524,6 +524,7 @@ class ResponseIntermediateStep(ResponseBaseModelIntermediate):
     type: str = "markdown"
     name: str
     payload: str
+    thought_text: str | None = None
 
 
 class ResponseObservabilityTrace(BaseModel, ResponseSerializable):

packages/nvidia_nat_core/src/nat/data_models/function.py

@@ -31,6 +31,8 @@ class FunctionBaseConfig(TypedBaseModel, BaseModelRegistryTag):
             If not provided, the function type will be used.
         `middleware`: List of function middleware names to apply to this function.
             These must match names defined in the `middleware` section of the YAML configuration.
+        `thought_description`: Optional custom thought description shown in the simplified thought
+            process UI in place of the default "Running function: <name>" text.
     """
     name: str | None = Field(
         default=None,
@@ -40,6 +42,11 @@ class FunctionBaseConfig(TypedBaseModel, BaseModelRegistryTag):
         default_factory=list,
         description="List of function middleware names to apply to this function in order",
     )
+    thought_description: str | None = Field(
+        default=None,
+        description="Optional custom thought description shown in the simplified thought process UI "
+        "in place of the default text for this function.",
+    )
 
 
 class FunctionGroupBaseConfig(TypedBaseModel, BaseModelRegistryTag):

packages/nvidia_nat_core/src/nat/front_ends/fastapi/step_adaptor.py

@@ -15,6 +15,7 @@
 
 import html
 import logging
+import typing
 from functools import reduce
 from textwrap import dedent
 
@@ -39,6 +40,17 @@ def __init__(self, config: StepAdaptorConfig):
         self._history: list[IntermediateStep] = []
         self.config = config
 
+    @staticmethod
+    def _get_thought_description(metadata: dict[str, typing.Any] | None, default: str, suffix: str = "") -> str:
+        """
+        Returns the user-configured `thought_description` from step metadata if present, otherwise falls back to
+        `default`. `suffix` is appended in either case (e.g. "..." while running, "... completed" once finished).
+        """
+        if not metadata or not metadata.get("thought_description"):
+            return default + suffix
+
+        return metadata["thought_description"] + suffix
+
     def _step_matches_filter(self, step: IntermediateStep, config: StepAdaptorConfig) -> bool:
         """
         Returns True if this intermediate step should be included (based on the config.mode).
@@ -203,10 +215,16 @@ def _handle_function(self, step: IntermediateStepPayload, ancestry: InvocationNo
             ```
             """).strip("\n").format(input_value=escaped_input, format_input_type=format_input_type)
 
+            # Omit thought_text for the top-level workflow function; its parent is the synthetic "root" node.
+            thought_text = None
+            if ancestry.parent_id != "root":
+                thought_text = self._get_thought_description(step.metadata, f"Running function: {step.name}", "...")
+
             event = ResponseIntermediateStep(id=step.UUID,
                                              name=f"Function Start: {step.name}",
                                              payload=payload_str,
-                                             parent_id=ancestry.parent_id)
+                                             parent_id=ancestry.parent_id,
+                                             thought_text=thought_text)
             return event
 
         if step.event_type == IntermediateStepType.FUNCTION_END:
@@ -256,10 +274,19 @@ def _handle_function(self, step: IntermediateStepPayload, ancestry: InvocationNo
                                     output_value=escaped_output,
                                     format_output_type=format_output_type)
 
+            # Omit thought_text for the top-level workflow function; its parent is the synthetic "root" node.
+            thought_text = None
+            if ancestry.parent_id != "root":
+                start_metadata = start_step.metadata if start_step else None
+                thought_text = self._get_thought_description(start_metadata,
+                                                             f"Running function: {step.name}",
+                                                             "... completed")
+
             event = ResponseIntermediateStep(id=step.UUID,
                                              name=f"Function Complete: {step.name}",
                                              payload=payload_str,
-                                             parent_id=ancestry.parent_id)
+                                             parent_id=ancestry.parent_id,
+                                             thought_text=thought_text)
             return event
 
         return None

packages/nvidia_nat_core/tests/nat/builder/test_function.py

@@ -21,12 +21,15 @@
 from pydantic import BaseModel
 
 from nat.builder.builder import Builder
+from nat.builder.context import Context
 from nat.builder.function import Function
 from nat.builder.function import LambdaFunction
 from nat.builder.function_info import FunctionInfo
 from nat.builder.workflow_builder import WorkflowBuilder
 from nat.cli.register_workflow import register_function
 from nat.data_models.function import FunctionBaseConfig
+from nat.data_models.intermediate_step import IntermediateStep
+from nat.data_models.intermediate_step import IntermediateStepType
 
 
 class DummyConfig(FunctionBaseConfig, name="dummy"):
@@ -90,6 +93,70 @@ async def test_direct_create_with_lambda():
         assert await fn_obj.ainvoke("test", to_type=str) == "test!"
 
 
+async def test_thought_description_included_in_function_start_metadata():
+    """
+    A function configured with `thought_description` should attach it as metadata on the
+    FUNCTION_START step it emits, so the step adaptor can surface it as a friendly thought label.
+    """
+    async with WorkflowBuilder() as builder:
+
+        fn_obj = await builder.add_function(name="test_function_with_thought",
+                                            config=LambdaFnConfig(thought_description="Searching the web"))
+
+        captured_steps: list[IntermediateStep] = []
+        Context.get().intermediate_step_manager.subscribe(captured_steps.append)
+
+        assert await fn_obj.ainvoke("test", to_type=str) == "test!"
+
+    start_steps = [s for s in captured_steps if s.event_type == IntermediateStepType.FUNCTION_START]
+
+    assert start_steps, "Expected at least one FUNCTION_START step to have been emitted"
+    assert start_steps[-1].metadata == {"thought_description": "Searching the web"}
+
+
+async def test_thought_description_absent_when_not_configured():
+    """
+    When `thought_description` is left unset (the default), no metadata should be attached to the
+    FUNCTION_START step, matching the pre-existing behavior for functions without a thought label.
+    """
+    async with WorkflowBuilder() as builder:
+
+        fn_obj = await builder.add_function(name="test_function_without_thought", config=LambdaFnConfig())
+
+        captured_steps: list[IntermediateStep] = []
+        Context.get().intermediate_step_manager.subscribe(captured_steps.append)
+
+        assert await fn_obj.ainvoke("test", to_type=str) == "test!"
+
+    start_steps = [s for s in captured_steps if s.event_type == IntermediateStepType.FUNCTION_START]
+
+    assert start_steps, "Expected at least one FUNCTION_START step to have been emitted"
+    assert start_steps[-1].metadata is None
+
+
+async def test_thought_description_included_in_function_start_metadata_for_streaming():
+    """
+    Same as `test_thought_description_included_in_function_start_metadata`, but for the `astream` path:
+    a function configured with `thought_description` should attach it as metadata on the FUNCTION_START
+    step it emits when invoked via streaming rather than `ainvoke`.
+    """
+    async with WorkflowBuilder() as builder:
+
+        fn_obj = await builder.add_function(name="test_stream_function_with_thought",
+                                            config=LambdaStreamFnConfig(thought_description="Searching the web"))
+
+        captured_steps: list[IntermediateStep] = []
+        Context.get().intermediate_step_manager.subscribe(captured_steps.append)
+
+        results = [result async for result in fn_obj.astream("test", to_type=str)]
+        assert results == ["test!"]
+
+    start_steps = [s for s in captured_steps if s.event_type == IntermediateStepType.FUNCTION_START]
+
+    assert start_steps, "Expected at least one FUNCTION_START step to have been emitted"
+    assert start_steps[-1].metadata == {"thought_description": "Searching the web"}
+
+
 async def test_direct_create_with_class():
 
     class ClassFnConfig(FunctionBaseConfig, name="test_class"):

packages/nvidia_nat_core/tests/nat/front_ends/fastapi/test_step_adaptor.py

@@ -78,17 +78,25 @@ def step_adaptor_disabled(disabled_config):
 def make_intermediate_step():
     """A factory fixture to create an IntermediateStep with minimal defaults."""
 
-    def _make_step(event_type: IntermediateStepType, data_input=None, data_output=None, name=None, UUID=None):
+    def _make_step(event_type: IntermediateStepType,
+                   data_input=None,
+                   data_output=None,
+                   name=None,
+                   UUID=None,
+                   metadata=None,
+                   ancestry_parent_id="abc"):
+        """Build a single IntermediateStep with the given event type, data, metadata, and ancestry."""
         payload = IntermediateStepPayload(
             event_type=event_type,
             name=name or "test_step",
             data=StreamEventData(input=data_input, output=data_output),
             UUID=UUID or "test-uuid-1234",
+            metadata=metadata,
         )
         # The IntermediateStep constructor requires a function_ancestry,
         # but for testing we can just pass None or a placeholder.
         return IntermediateStep(parent_id="root",
-                                function_ancestry=InvocationNode(parent_id="abc",
+                                function_ancestry=InvocationNode(parent_id=ancestry_parent_id,
                                                                  function_id="def",
                                                                  function_name="xyz"),
                                 payload=payload)
@@ -496,6 +504,123 @@ def test_process_function_end_without_output(step_adaptor_default, make_intermed
     assert step_adaptor_default._history[-1] is step
 
 
+def test_get_thought_description_falls_back_to_default():
+    """
+    `_get_thought_description` should return the default text when metadata is missing, empty, or has an
+    empty `thought_description` value.
+    """
+    result = StepAdaptor._get_thought_description(None, "Running function: foo", "...")
+    assert result == "Running function: foo..."
+
+    result = StepAdaptor._get_thought_description({}, "Running function: foo", "...")
+    assert result == "Running function: foo..."
+
+    result = StepAdaptor._get_thought_description({"thought_description": ""}, "Running function: foo", "...")
+    assert result == "Running function: foo..."
+
+
+def test_get_thought_description_uses_custom_value():
+    """
+    `_get_thought_description` should return the configured `thought_description` (plus suffix) when present.
+    """
+    result = StepAdaptor._get_thought_description({"thought_description": "Searching the web"},
+                                                  "Running function: foo",
+                                                  "...")
+    assert result == "Searching the web..."
+
+
+def test_function_start_uses_default_thought_text(step_adaptor_default, make_intermediate_step):
+    """
+    A FUNCTION_START event with no `thought_description` metadata should get the default thought text.
+    """
+    step = make_intermediate_step(
+        event_type=IntermediateStepType.FUNCTION_START,
+        data_input="Function Input Data",
+        name="test_function",
+    )
+
+    result = step_adaptor_default.process(step)
+
+    assert result.thought_text == "Running function: test_function..."
+
+
+def test_function_start_uses_custom_thought_description(step_adaptor_default, make_intermediate_step):
+    """
+    A FUNCTION_START event whose function was configured with `thought_description` should surface that
+    text (with the "..." suffix) instead of the default "Running function: <name>" text.
+    """
+    step = make_intermediate_step(
+        event_type=IntermediateStepType.FUNCTION_START,
+        data_input="Function Input Data",
+        name="test_function",
+        metadata={"thought_description": "Searching the web"},
+    )
+
+    result = step_adaptor_default.process(step)
+
+    assert result.thought_text == "Searching the web..."
+
+
+def test_function_start_omits_thought_text_for_root_workflow_function(step_adaptor_default, make_intermediate_step):
+    """
+    The top-level workflow function's parent is the synthetic "root" node; it should not get a thought_text
+    since it represents the whole run rather than a sub-step.
+    """
+    step = make_intermediate_step(
+        event_type=IntermediateStepType.FUNCTION_START,
+        data_input="Function Input Data",
+        name="test_workflow",
+        ancestry_parent_id="root",
+    )
+
+    result = step_adaptor_default.process(step)
+
+    assert result.thought_text is None
+
+
+def test_function_end_uses_custom_thought_description_from_start_event(step_adaptor_default, make_intermediate_step):
+    """
+    FUNCTION_END should read `thought_description` from the matching FUNCTION_START event's metadata (the
+    end payload itself carries no metadata) and append the "completed" suffix.
+    """
+    uuid = "function-thought-uuid"
+    start_step = make_intermediate_step(
+        event_type=IntermediateStepType.FUNCTION_START,
+        data_input="Function Input Data",
+        name="test_function",
+        UUID=uuid,
+        metadata={"thought_description": "Searching the web"},
+    )
+    end_step = make_intermediate_step(
+        event_type=IntermediateStepType.FUNCTION_END,
+        data_output="Function Output Data",
+        name="test_function",
+        UUID=uuid,
+    )
+
+    step_adaptor_default.process(start_step)
+    result = step_adaptor_default.process(end_step)
+
+    assert result.thought_text == "Searching the web... completed"
+
+
+def test_function_end_omits_thought_text_for_root_workflow_function(step_adaptor_default, make_intermediate_step):
+    """
+    The top-level workflow function's parent is the synthetic "root" node; FUNCTION_END should not get a
+    thought_text for it either, mirroring the FUNCTION_START behavior.
+    """
+    step = make_intermediate_step(
+        event_type=IntermediateStepType.FUNCTION_END,
+        data_output="Function Output Data",
+        name="test_workflow",
+        ancestry_parent_id="root",
+    )
+
+    result = step_adaptor_default.process(step)
+
+    assert result.thought_text is None
+
+
 def test_function_events_in_custom_mode(step_adaptor_custom, make_intermediate_step):
     """
     In CUSTOM mode without FUNCTION_START/END in custom_event_types,