feat(sdk): chained entity path on nested tasks (#1782)

Co-authored-by: Nir Gazit <nirga@users.noreply.github.com>
traceloop · Aug 7, 2024 · 826a101 · 826a101
1 parent 39f9db6
commit 826a101
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 26 deletions.
diff --git a/packages/traceloop-sdk/tests/test_nested_tasks.py b/packages/traceloop-sdk/tests/test_nested_tasks.py
@@ -1,5 +1,6 @@
 from opentelemetry.semconv_ai import SpanAttributes
 from traceloop.sdk.decorators import task, workflow
+from pytest import raises
 
 
 def test_nested_tasks(exporter):
@@ -32,14 +33,15 @@ def inner_inner_task():
     inner_inner_task_span = spans[0]
     inner_task_span = spans[1]
     outer_task_span = spans[2]
+    some_workflow_span = spans[3]
+
     assert (
-        inner_inner_task_span.attributes[SpanAttributes.TRACELOOP_ENTITY_NAME]
-        == "outer_task.inner_task.inner_inner_task"
-    )
-    assert (
-        inner_task_span.attributes[SpanAttributes.TRACELOOP_ENTITY_NAME]
-        == "outer_task.inner_task"
+        inner_inner_task_span.attributes[SpanAttributes.TRACELOOP_ENTITY_PATH] ==
+        "outer_task.inner_task"
     )
     assert (
-        outer_task_span.attributes[SpanAttributes.TRACELOOP_ENTITY_NAME] == "outer_task"
+        inner_task_span.attributes[SpanAttributes.TRACELOOP_ENTITY_PATH] == "outer_task"
     )
+    with raises(KeyError):
+        _ = outer_task_span.attributes[SpanAttributes.TRACELOOP_ENTITY_PATH]
+        _ = some_workflow_span.attributes[SpanAttributes.TRACELOOP_ENTITY_PATH]
diff --git a/packages/traceloop-sdk/traceloop/sdk/decorators/base.py b/packages/traceloop-sdk/traceloop/sdk/decorators/base.py
@@ -12,8 +12,8 @@
 from traceloop.sdk.tracing import get_tracer, set_workflow_name
 from traceloop.sdk.tracing.tracing import (
     TracerWrapper,
-    set_entity_name,
-    get_chained_entity_name,
+    set_entity_path,
+    get_chained_entity_path,
 )
 from traceloop.sdk.utils import camel_to_snake
 from traceloop.sdk.utils.json_encoder import JSONEncoder
@@ -47,16 +47,14 @@ def wrap(*args, **kwargs):
                     TraceloopSpanKindValues.TASK,
                     TraceloopSpanKindValues.TOOL,
                 ]:
-                    chained_entity_name = get_chained_entity_name(entity_name)
-                    set_entity_name(chained_entity_name)
-                else:
-                    chained_entity_name = entity_name
+                    entity_path = get_chained_entity_path(entity_name)
+                    set_entity_path(entity_path)
 
                 span.set_attribute(
                     SpanAttributes.TRACELOOP_SPAN_KIND, tlp_span_kind.value
                 )
                 span.set_attribute(
-                    SpanAttributes.TRACELOOP_ENTITY_NAME, chained_entity_name
+                    SpanAttributes.TRACELOOP_ENTITY_NAME, entity_name
                 )
                 if version:
                     span.set_attribute(SpanAttributes.TRACELOOP_ENTITY_VERSION, version)
@@ -149,16 +147,14 @@ async def wrap(*args, **kwargs):
                     TraceloopSpanKindValues.TASK,
                     TraceloopSpanKindValues.TOOL,
                 ]:
-                    chained_entity_name = get_chained_entity_name(entity_name)
-                    set_entity_name(chained_entity_name)
-                else:
-                    chained_entity_name = entity_name
+                    entity_path = get_chained_entity_path(entity_name)
+                    set_entity_path(entity_path)
 
                 span.set_attribute(
                     SpanAttributes.TRACELOOP_SPAN_KIND, tlp_span_kind.value
                 )
                 span.set_attribute(
-                    SpanAttributes.TRACELOOP_ENTITY_NAME, chained_entity_name
+                    SpanAttributes.TRACELOOP_ENTITY_NAME, entity_name
                 )
                 if version:
                     span.set_attribute(SpanAttributes.TRACELOOP_ENTITY_VERSION, version)

diff --git a/packages/traceloop-sdk/traceloop/sdk/tracing/tracing.py b/packages/traceloop-sdk/traceloop/sdk/tracing/tracing.py
@@ -329,9 +329,9 @@ def _span_processor_on_start(self, span, parent_context):
         if workflow_name is not None:
             span.set_attribute(SpanAttributes.TRACELOOP_WORKFLOW_NAME, workflow_name)
 
-        entity_name = get_value("entity_name")
-        if entity_name is not None:
-            span.set_attribute(SpanAttributes.TRACELOOP_ENTITY_NAME, entity_name)
+        entity_path = get_value("entity_path")
+        if entity_path is not None:
+            span.set_attribute(SpanAttributes.TRACELOOP_ENTITY_PATH, entity_path)
 
         association_properties = get_value("association_properties")
         if association_properties is not None:
@@ -444,12 +444,12 @@ def set_workflow_name(workflow_name: str) -> None:
     attach(set_value("workflow_name", workflow_name))
 
 
-def set_entity_name(entity_name: str) -> None:
-    attach(set_value("entity_name", entity_name))
+def set_entity_path(entity_path: str) -> None:
+    attach(set_value("entity_path", entity_path))
 
 
-def get_chained_entity_name(entity_name: str) -> str:
-    parent = get_value("entity_name")
+def get_chained_entity_path(entity_name: str) -> str:
+    parent = get_value("entity_path")
     if parent is None:
         return entity_name
     else: