Record root span name for each child span

src/promptflow-azure/promptflow/azure/_storage/cosmosdb/summary.py

@@ -99,17 +99,18 @@ def _persist_running_item(self, client: ContainerProxy):
         trace_id = self.span.trace_id
         session_id = self.session_id
 
+        attributes: dict = self.span.attributes
         item = SummaryLine(
             id=trace_id,
             partition_key=self.collection_id,
             session_id=session_id,
             trace_id=trace_id,
+            name=attributes.get(SpanAttributeFieldName.ROOT_SPAN_NAME, None),
             status=RUNNING_LINE_RUN_STATUS,
             collection_id=self.collection_id,
             created_by=self.created_by,
             start_time=self.span.start_time.isoformat(),
         )
-        attributes: dict = self.span.attributes
         if SpanAttributeFieldName.LINE_RUN_ID in attributes:
             item.line_run_id = attributes[SpanAttributeFieldName.LINE_RUN_ID]
         elif SpanAttributeFieldName.BATCH_RUN_ID in attributes and SpanAttributeFieldName.LINE_NUMBER in attributes:

src/promptflow-core/promptflow/_constants.py

@@ -154,6 +154,7 @@ class SpanAttributeFieldName:
     FUNCTION = "function"
     INPUTS = "inputs"
     OUTPUT = "output"
+    ROOT_SPAN_NAME = "root_span_name"  # Record root span name in all child spans to show correct name for running item
     # token metrics
     COMPLETION_TOKEN_COUNT = "llm.usage.completion_tokens"
     PROMPT_TOKEN_COUNT = "llm.usage.prompt_tokens"

src/promptflow-core/promptflow/executor/flow_executor.py

@@ -843,6 +843,15 @@ def _start_flow_span(self, inputs: Mapping[str, Any]):
         with otel_tracer.start_as_current_span(self._flow.name) as span:
             # Store otel trace id in context for correlation
             OperationContext.get_instance()["otel_trace_id"] = f"0x{format_trace_id(span.get_span_context().trace_id)}"
+
+            root_span_name_key = "root_span_name"
+            attributes = OperationContext.get_instance()._get_otel_attributes()
+            exist_root_span_name = root_span_name_key in attributes
+            # Add root span name to attributes to make sure we could show correct name for running flow in UI.
+            # For scenario like chat group, the root span name is already set, we don't need to set it again.
+            if not exist_root_span_name:
+                OperationContext.get_instance()._add_otel_attributes(root_span_name_key, span.name)
+
             # initialize span
             span.set_attributes(
                 {
@@ -854,6 +863,8 @@ def _start_flow_span(self, inputs: Mapping[str, Any]):
             # enrich span with input
             enrich_span_with_input(span, inputs)
             yield span
+            if not exist_root_span_name:
+                OperationContext.get_instance()._remove_otel_attributes(root_span_name_key)
 
     async def _exec_inner_with_trace_async(
         self,

src/promptflow-tracing/promptflow/tracing/_trace.py

@@ -39,6 +39,18 @@ def _record_keyboard_interrupt_to_span(span: Span):
         raise
 
 
+@contextlib.contextmanager
+def _handle_root_span_name(span: Span):
+    root_span_name_key = "root_span_name"
+    attributes = OperationContext.get_instance()._get_otel_attributes()
+    exist_root_span_name = root_span_name_key in attributes
+    if not exist_root_span_name:
+        OperationContext.get_instance()._add_otel_attributes(root_span_name_key, span.name)
+    yield
+    if not exist_root_span_name:
+        OperationContext.get_instance()._remove_otel_attributes(root_span_name_key)
+
+
 class TokenCollector:
     _lock = Lock()
 
@@ -187,7 +199,7 @@ def traced_generator(original_span: ReadableSpan, inputs, generator):
     with otel_tracer.start_as_current_span(
         f"Iterated({original_span.name})",
         links=[link],
-    ) as span, _record_keyboard_interrupt_to_span(span):
+    ) as span, _handle_root_span_name(span), _record_keyboard_interrupt_to_span(span):
         enrich_span_with_original_attributes(span, original_span.attributes)
         # Enrich the new span with input before generator iteration to prevent loss of input information.
         # The input is as an event within this span.
@@ -211,7 +223,7 @@ async def traced_async_generator(original_span: ReadableSpan, inputs, generator)
     with otel_tracer.start_as_current_span(
         f"Iterated({original_span.name})",
         links=[link],
-    ) as span, _record_keyboard_interrupt_to_span(span):
+    ) as span, _handle_root_span_name(span), _record_keyboard_interrupt_to_span(span):
         enrich_span_with_original_attributes(span, original_span.attributes)
         # Enrich the new span with input before generator iteration to prevent loss of input information.
         # The input is as an event within this span.
@@ -388,7 +400,9 @@ async def wrapped(*args, **kwargs):
         span_name = get_node_name_from_context(used_for_span_name=True) or trace.name
         # need to get everytime to ensure tracer is latest
         otel_tracer = otel_trace.get_tracer("promptflow")
-        with otel_tracer.start_as_current_span(span_name) as span, _record_keyboard_interrupt_to_span(span):
+        with otel_tracer.start_as_current_span(span_name) as span, _handle_root_span_name(
+            span
+        ), _record_keyboard_interrupt_to_span(span):
             # Store otel trace id in context for correlation
             OperationContext.get_instance()["otel_trace_id"] = f"0x{format_trace_id(span.get_span_context().trace_id)}"
             enrich_span_with_trace(span, trace)
@@ -454,7 +468,9 @@ def wrapped(*args, **kwargs):
         span_name = get_node_name_from_context(used_for_span_name=True) or trace.name
         # need to get everytime to ensure tracer is latest
         otel_tracer = otel_trace.get_tracer("promptflow")
-        with otel_tracer.start_as_current_span(span_name) as span, _record_keyboard_interrupt_to_span(span):
+        with otel_tracer.start_as_current_span(span_name) as span, _handle_root_span_name(
+            span
+        ), _record_keyboard_interrupt_to_span(span):
             # Store otel trace id in context for correlation
             OperationContext.get_instance()["otel_trace_id"] = f"0x{format_trace_id(span.get_span_context().trace_id)}"
             enrich_span_with_trace(span, trace)