diff --git a/packages/nvidia_nat_core/src/nat/builder/function.py b/packages/nvidia_nat_core/src/nat/builder/function.py index 23a3457892..5443e3384b 100644 --- a/packages/nvidia_nat_core/src/nat/builder/function.py +++ b/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) diff --git a/packages/nvidia_nat_core/src/nat/data_models/api_server.py b/packages/nvidia_nat_core/src/nat/data_models/api_server.py index 37bdf39ff2..d557835e17 100644 --- a/packages/nvidia_nat_core/src/nat/data_models/api_server.py +++ b/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): diff --git a/packages/nvidia_nat_core/src/nat/data_models/function.py b/packages/nvidia_nat_core/src/nat/data_models/function.py index a089682a97..40a73a00ef 100644 --- a/packages/nvidia_nat_core/src/nat/data_models/function.py +++ b/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: " 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): 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..3d6e2d03e3 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,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 diff --git a/packages/nvidia_nat_core/tests/nat/builder/test_function.py b/packages/nvidia_nat_core/tests/nat/builder/test_function.py index a4cc8f6caa..c99f2f45bd 100644 --- a/packages/nvidia_nat_core/tests/nat/builder/test_function.py +++ b/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"): 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..4eb21246b7 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 @@ -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: " 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,