Skip to content

Commit 1925cb8

Browse files
thepatrickchinthepatrickchin
committed
fix: update thought emission to use data field instead of metadata
Ensures thought text is properly displayed in observability traces Signed-off-by: Patrick Chin <8509935+thepatrickchin@users.noreply.github.com>
1 parent 4fe4a27 commit 1925cb8

File tree

4 files changed

+34
-40
lines changed

4 files changed

+34
-40
lines changed

‎packages/nvidia_nat_core/src/nat/builder/thought.py‎

Lines changed: 5 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -49,16 +49,14 @@ def emit_thought(context: "Context", thought_text: str, name: str | None = None)
4949
IntermediateStepPayload(UUID=thought_uuid,
5050
event_type=IntermediateStepType.SPAN_START,
5151
name=thought_name,
52-
data=StreamEventData(input=None),
53-
metadata={"thought_text": thought_text}))
52+
data=StreamEventData(input=thought_text)))
5453

5554
# Immediately emit END event (complete thought)
5655
context.intermediate_step_manager.push_intermediate_step(
5756
IntermediateStepPayload(UUID=thought_uuid,
5857
event_type=IntermediateStepType.SPAN_END,
5958
name=thought_name,
60-
data=StreamEventData(output=None),
61-
metadata={"thought_text": thought_text}))
59+
data=StreamEventData(output=thought_text)))
6260

6361
return thought_uuid
6462

@@ -84,8 +82,7 @@ def emit_thought_start(context: "Context", thought_text: str, name: str | None =
8482
IntermediateStepPayload(UUID=thought_uuid,
8583
event_type=IntermediateStepType.SPAN_START,
8684
name=thought_name,
87-
data=StreamEventData(input=None),
88-
metadata={"thought_text": thought_text}))
85+
data=StreamEventData(input=thought_text)))
8986

9087
return thought_uuid
9188

@@ -104,8 +101,7 @@ def emit_thought_chunk(context: "Context", thought_uuid: str, thought_text: str)
104101
IntermediateStepPayload(UUID=thought_uuid,
105102
event_type=IntermediateStepType.SPAN_CHUNK,
106103
name="custom_thought",
107-
data=StreamEventData(chunk=thought_text),
108-
metadata={"thought_text": thought_text}))
104+
data=StreamEventData(chunk=thought_text)))
109105

110106

111107
def emit_thought_end(context: "Context", thought_uuid: str, thought_text: str | None = None) -> None:
@@ -120,5 +116,4 @@ def emit_thought_end(context: "Context", thought_uuid: str, thought_text: str |
120116
IntermediateStepPayload(UUID=thought_uuid,
121117
event_type=IntermediateStepType.SPAN_END,
122118
name="custom_thought",
123-
data=StreamEventData(output=None),
124-
metadata={"thought_text": thought_text} if thought_text else {}))
119+
data=StreamEventData(output=thought_text)))

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

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -328,11 +328,15 @@ def _handle_span(self, step: IntermediateStepPayload, ancestry: InvocationNode)
328328
"""
329329
Handles SPAN events (SPAN_START, SPAN_CHUNK, SPAN_END) for custom thoughts.
330330
"""
331-
# Check if this is a custom thought with metadata
332-
if not step.metadata or not isinstance(step.metadata, dict):
333-
return None
331+
# Extract thought_text from data field based on event type
332+
thought_text = None
333+
if step.event_type == IntermediateStepType.SPAN_START and step.data:
334+
thought_text = step.data.input
335+
elif step.event_type == IntermediateStepType.SPAN_CHUNK and step.data:
336+
thought_text = step.data.chunk
337+
elif step.event_type == IntermediateStepType.SPAN_END and step.data:
338+
thought_text = step.data.output
334339

335-
thought_text = step.metadata.get("thought_text")
336340
if not thought_text:
337341
return None
338342

@@ -341,7 +345,7 @@ def _handle_span(self, step: IntermediateStepPayload, ancestry: InvocationNode)
341345
name=step.name or "Span",
342346
payload="",
343347
parent_id=ancestry.function_id,
344-
thought_text=thought_text)
348+
thought_text=str(thought_text))
345349

346350
def _handle_custom(self, payload: IntermediateStepPayload, ancestry: InvocationNode) -> ResponseSerializable | None:
347351
"""

‎packages/nvidia_nat_core/tests/nat/builder/test_thought.py‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -98,9 +98,9 @@ def test_emit_thought_creates_start_and_end_events(ctx: Context, output_steps: l
9898
assert start_event.payload.event_type == IntermediateStepType.SPAN_START
9999
assert end_event.payload.event_type == IntermediateStepType.SPAN_END
100100

101-
# Check thought_text in metadata
102-
assert start_event.payload.metadata["thought_text"] == "Processing data..."
103-
assert end_event.payload.metadata["thought_text"] == "Processing data..."
101+
# Check thought_text in data
102+
assert start_event.payload.data.input == "Processing data..."
103+
assert end_event.payload.data.output == "Processing data..."
104104

105105

106106
def test_emit_thought_custom_name(ctx: Context, output_steps: list[IntermediateStep]):
@@ -126,7 +126,7 @@ def test_emit_thought_start_creates_start_event(ctx: Context, output_steps: list
126126
start_event = output_steps[0]
127127
assert start_event.payload.UUID == thought_uuid
128128
assert start_event.payload.event_type == IntermediateStepType.SPAN_START
129-
assert start_event.payload.metadata["thought_text"] == "Starting process..."
129+
assert start_event.payload.data.input == "Starting process..."
130130

131131
# Balance span stack for fixture teardown
132132
emit_thought_end(ctx, thought_uuid)
@@ -150,7 +150,7 @@ def test_emit_thought_chunk_creates_chunk_event(ctx: Context, output_steps: list
150150
chunk_event = output_steps[0]
151151
assert chunk_event.payload.UUID == thought_uuid
152152
assert chunk_event.payload.event_type == IntermediateStepType.SPAN_CHUNK
153-
assert chunk_event.payload.metadata["thought_text"] == "Processing: 50%"
153+
assert chunk_event.payload.data.chunk == "Processing: 50%"
154154

155155
emit_thought_end(ctx, thought_uuid)
156156

@@ -173,7 +173,7 @@ def test_emit_thought_end_creates_end_event(ctx: Context, output_steps: list[Int
173173
end_event = output_steps[0]
174174
assert end_event.payload.UUID == thought_uuid
175175
assert end_event.payload.event_type == IntermediateStepType.SPAN_END
176-
assert end_event.payload.metadata["thought_text"] == "Complete"
176+
assert end_event.payload.data.output == "Complete"
177177

178178

179179
def test_emit_thought_end_with_no_text(ctx: Context, output_steps: list[IntermediateStep]):
@@ -186,8 +186,8 @@ def test_emit_thought_end_with_no_text(ctx: Context, output_steps: list[Intermed
186186
end_event = output_steps[0]
187187
assert end_event.payload.UUID == thought_uuid
188188
assert end_event.payload.event_type == IntermediateStepType.SPAN_END
189-
# metadata should be empty dict when no thought_text provided
190-
assert end_event.payload.metadata == {}
189+
# data.output should be None when no thought_text provided
190+
assert end_event.payload.data.output is None
191191

192192

193193
# --------------------------------------------------------------------------- #

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

Lines changed: 12 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -813,13 +813,12 @@ def test_tool_end_with_thought_description_override(step_adaptor_default, make_i
813813
# Tests for custom thought emission via SPAN events
814814
# --------------------
815815
def test_span_start_with_custom_thought(step_adaptor_default):
816-
"""Test that SPAN_START events with thought_text metadata are processed correctly."""
816+
"""Test that SPAN_START events with thought_text in data.input are processed correctly."""
817817
custom_thought = "Processing data: initializing..."
818818
payload = IntermediateStepPayload(event_type=IntermediateStepType.SPAN_START,
819819
name="custom_thought",
820-
data=StreamEventData(input=None),
821-
UUID="span-uuid-1",
822-
metadata={"thought_text": custom_thought})
820+
data=StreamEventData(input=custom_thought),
821+
UUID="span-uuid-1")
823822
step = IntermediateStep(parent_id="root",
824823
function_ancestry=InvocationNode(parent_id="abc", function_id="def", function_name="xyz"),
825824
payload=payload)
@@ -838,9 +837,8 @@ def test_span_chunk_with_custom_thought(step_adaptor_default):
838837

839838
payload_start = IntermediateStepPayload(event_type=IntermediateStepType.SPAN_START,
840839
name="custom_thought",
841-
data=StreamEventData(input=None),
842-
UUID=uuid,
843-
metadata={"thought_text": "Processing: 0%"})
840+
data=StreamEventData(input="Processing: 0%"),
841+
UUID=uuid)
844842
step_start = IntermediateStep(parent_id="root",
845843
function_ancestry=InvocationNode(parent_id="abc",
846844
function_id="def",
@@ -851,9 +849,8 @@ def test_span_chunk_with_custom_thought(step_adaptor_default):
851849
custom_thought_update = "Processing: 50%"
852850
payload_chunk = IntermediateStepPayload(event_type=IntermediateStepType.SPAN_CHUNK,
853851
name="custom_thought",
854-
data=StreamEventData(chunk="50%"),
855-
UUID=uuid,
856-
metadata={"thought_text": custom_thought_update})
852+
data=StreamEventData(chunk=custom_thought_update),
853+
UUID=uuid)
857854
step_chunk = IntermediateStep(parent_id="root",
858855
function_ancestry=InvocationNode(parent_id="abc",
859856
function_id="def",
@@ -867,13 +864,11 @@ def test_span_chunk_with_custom_thought(step_adaptor_default):
867864

868865

869866
def test_span_without_thought_text_returns_none(step_adaptor_default):
870-
"""Test that SPAN events without thought_text metadata return None."""
871-
payload = IntermediateStepPayload(
872-
event_type=IntermediateStepType.SPAN_START,
873-
name="regular_span",
874-
data=StreamEventData(input=None),
875-
UUID="span-uuid-no-thought",
876-
)
867+
"""Test that SPAN events without thought_text in data return None."""
868+
payload = IntermediateStepPayload(event_type=IntermediateStepType.SPAN_START,
869+
name="regular_span",
870+
data=StreamEventData(input=None),
871+
UUID="span-uuid-no-thought")
877872
step = IntermediateStep(parent_id="root",
878873
function_ancestry=InvocationNode(parent_id="abc", function_id="def", function_name="xyz"),
879874
payload=payload)

0 commit comments

Comments
 (0)