Code review - use parent_run_id to demote from workflow to task

traceloop · Jun 29, 2024 · 27024dd · 27024dd
1 parent 4f7d45e
commit 27024dd
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 16 deletions.
diff --git a/...try-instrumentation-langchain/opentelemetry/instrumentation/langchain/callback_wrapper.py b/...try-instrumentation-langchain/opentelemetry/instrumentation/langchain/callback_wrapper.py
@@ -6,7 +6,7 @@
 from langchain_core.callbacks import BaseCallbackHandler, BaseCallbackManager
 from langchain_core.messages import BaseMessage
 from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY
-from opentelemetry.semconv.ai import LLMRequestTypeValues, SpanAttributes
+from opentelemetry.semconv.ai import LLMRequestTypeValues, SpanAttributes, TraceloopSpanKindValues
 from opentelemetry.context.context import Context
 from opentelemetry.trace import set_span_in_context, Tracer
 from opentelemetry.trace.span import Span
@@ -40,11 +40,11 @@ def from_args(cls, kind: str, instance_name: str, args: Any) -> "Instance":
             run_name = None
         return cls(kind, instance_name, run_name)
 
-    def get_kind(self) -> str:
-        return self.kind
-
-    def get_name(self) -> str:
-        return f"{self.run_name if self.run_name is not None else self.class_name}.langchain.{self.kind}"
+    def get_name(self, kind: Optional[str] = None) -> str:
+        return (
+            f"{self.run_name if self.run_name is not None else self.class_name}"
+            f".langchain.{self.kind if kind is None else kind}"
+        )
 
 
 @dataclass
@@ -133,14 +133,14 @@ def _end_span(self, span: Span, run_id: UUID) -> None:
 
     def _create_span(self, run_id: UUID, parent_run_id: Optional[UUID], class_name: str) -> Span:
         instance = self.instances[class_name]
-        kind = instance.get_kind()
-        name = instance.get_name()
+        kind = instance.kind if parent_run_id is None else TraceloopSpanKindValues.TASK.value
+        name = instance.get_name(kind)
         if parent_run_id is not None:
             span = self.tracer.start_span(name, context=self.spans[parent_run_id].context)
         else:
             span = self.tracer.start_span(name)
         span.set_attribute(SpanAttributes.TRACELOOP_SPAN_KIND, kind)
-        span.set_attribute(SpanAttributes.TRACELOOP_ENTITY_NAME, f"{name}.langchain.{kind}")
+        span.set_attribute(SpanAttributes.TRACELOOP_ENTITY_NAME, name)
         current_context = set_span_in_context(span)
         token = context_api.attach(current_context)
         self.spans[run_id] = SpanHolder(span, token, current_context, [])

diff --git a/packages/opentelemetry-instrumentation-langchain/tests/test_chains.py b/packages/opentelemetry-instrumentation-langchain/tests/test_chains.py
@@ -45,13 +45,13 @@ def test_sequential_chain(exporter):
 
     assert [
         "openai.completion",
-        "LLMChain.langchain.workflow",
+        "LLMChain.langchain.task",
         "openai.completion",
-        "LLMChain.langchain.workflow",
+        "LLMChain.langchain.task",
         "SequentialChain.langchain.workflow",
     ] == [span.name for span in spans]
 
-    synopsis_span, review_span = [span for span in spans if span.name == "LLMChain.langchain.workflow"]
+    synopsis_span, review_span = [span for span in spans if span.name == "LLMChain.langchain.task"]
 
     data = json.loads(synopsis_span.attributes[SpanAttributes.TRACELOOP_ENTITY_INPUT])
     assert data["inputs"] == {'title': 'Tragedy at sunset on the beach', 'era': 'Victorian England'}
@@ -112,13 +112,13 @@ async def test_asequential_chain(exporter):
 
     assert [
         "openai.completion",
-        "LLMChain.langchain.workflow",
+        "LLMChain.langchain.task",
         "openai.completion",
-        "LLMChain.langchain.workflow",
+        "LLMChain.langchain.task",
         "SequentialChain.langchain.workflow",
     ] == [span.name for span in spans]
 
-    synopsis_span, review_span = [span for span in spans if span.name == "LLMChain.langchain.workflow"]
+    synopsis_span, review_span = [span for span in spans if span.name == "LLMChain.langchain.task"]
 
     data = json.loads(synopsis_span.attributes[SpanAttributes.TRACELOOP_ENTITY_INPUT])
     assert data["inputs"] == {'title': 'Tragedy at sunset on the beach', 'era': 'Victorian England'}

diff --git a/packages/opentelemetry-instrumentation-langchain/tests/test_documents_chains.py b/packages/opentelemetry-instrumentation-langchain/tests/test_documents_chains.py
@@ -41,7 +41,7 @@ def test_sequential_chain(exporter):
     spans = exporter.get_finished_spans()
 
     assert [
-        "LLMChain.langchain.workflow",
+        "LLMChain.langchain.task",
         "StuffDocumentsChain.langchain.workflow",
     ] == [span.name for span in spans]