Reduce duplicate code

Author: sjrl

haystack/core/pipeline/async_pipeline.py

@@ -11,10 +11,8 @@
 from haystack.core.errors import PipelineMaxComponentRuns, PipelineRuntimeError
 from haystack.core.pipeline.base import (
     _COMPONENT_INPUT,
-    _COMPONENT_NAME,
     _COMPONENT_OUTPUT,
-    _COMPONENT_RUN,
-    _COMPONENT_TYPE,
+    _COMPONENT_VISITS,
     ComponentPriority,
     PipelineBase,
 )
@@ -57,28 +55,8 @@ async def _run_component_async(  # pylint: disable=too-many-positional-arguments
             raise PipelineMaxComponentRuns(f"Max runs for '{component_name}' reached.")
 
         instance: Component = component["instance"]
-        with tracing.tracer.trace(
-            _COMPONENT_RUN,
-            tags={
-                _COMPONENT_NAME: component_name,
-                _COMPONENT_TYPE: instance.__class__.__name__,
-                "haystack.component.input_types": {k: type(v).__name__ for k, v in component_inputs.items()},
-                "haystack.component.input_spec": {
-                    key: {
-                        "type": (value.type.__name__ if isinstance(value.type, type) else str(value.type)),
-                        "senders": value.senders,
-                    }
-                    for key, value in instance.__haystack_input__._sockets_dict.items()  # type: ignore
-                },
-                "haystack.component.output_spec": {
-                    key: {
-                        "type": (value.type.__name__ if isinstance(value.type, type) else str(value.type)),
-                        "receivers": value.receivers,
-                    }
-                    for key, value in instance.__haystack_output__._sockets_dict.items()  # type: ignore
-                },
-            },
-            parent_span=parent_span,
+        with PipelineBase._create_component_span(
+            component_name=component_name, instance=instance, inputs=component_inputs, parent_span=parent_span
         ) as span:
             span.set_content_tag(_COMPONENT_INPUT, deepcopy(component_inputs))
             logger.info("Running component {component_name}", component_name=component_name)
@@ -97,7 +75,7 @@ async def _run_component_async(  # pylint: disable=too-many-positional-arguments
             if not isinstance(outputs, dict):
                 raise PipelineRuntimeError.from_invalid_output(component_name, instance.__class__, outputs)
 
-            span.set_tag("haystack.component.visits", component_visits[component_name])
+            span.set_tag(_COMPONENT_VISITS, component_visits[component_name])
             span.set_content_tag(_COMPONENT_OUTPUT, deepcopy(outputs))
 
             return outputs

haystack/core/pipeline/base.py

@@ -12,7 +12,7 @@
 
 import networkx  # type:ignore
 
-from haystack import logging
+from haystack import logging, tracing
 from haystack.core.component import Component, InputSocket, OutputSocket, component
 from haystack.core.errors import (
     DeserializationError,
@@ -54,11 +54,9 @@
 
 
 # Constants for tracing tags
-_COMPONENT_RUN = "haystack.component.run"
-_COMPONENT_NAME = "haystack.component.name"
-_COMPONENT_TYPE = "haystack.component.type"
 _COMPONENT_INPUT = "haystack.component.input"
 _COMPONENT_OUTPUT = "haystack.component.output"
+_COMPONENT_VISITS = "haystack.component.visits"
 
 
 class ComponentPriority(IntEnum):
@@ -776,6 +774,34 @@ def warm_up(self):
                 logger.info("Warming up component {node}...", node=node)
                 self.graph.nodes[node]["instance"].warm_up()
 
+    @staticmethod
+    def _create_component_span(
+        component_name: str, instance: Component, inputs: Dict[str, Any], parent_span: Optional[tracing.Span] = None
+    ):
+        return tracing.tracer.trace(
+            "haystack.component.run",
+            tags={
+                "haystack.component.name": component_name,
+                "haystack.component.type": instance.__class__.__name__,
+                "haystack.component.input_types": {k: type(v).__name__ for k, v in inputs.items()},
+                "haystack.component.input_spec": {
+                    key: {
+                        "type": (value.type.__name__ if isinstance(value.type, type) else str(value.type)),
+                        "senders": value.senders,
+                    }
+                    for key, value in instance.__haystack_input__._sockets_dict.items()  # type: ignore
+                },
+                "haystack.component.output_spec": {
+                    key: {
+                        "type": (value.type.__name__ if isinstance(value.type, type) else str(value.type)),
+                        "receivers": value.receivers,
+                    }
+                    for key, value in instance.__haystack_output__._sockets_dict.items()  # type: ignore
+                },
+            },
+            parent_span=parent_span,
+        )
+
     def _validate_input(self, data: Dict[str, Any]):
         """
         Validates pipeline input data.

haystack/core/pipeline/pipeline.py

@@ -10,10 +10,8 @@
 from haystack.core.errors import PipelineRuntimeError
 from haystack.core.pipeline.base import (
     _COMPONENT_INPUT,
-    _COMPONENT_NAME,
     _COMPONENT_OUTPUT,
-    _COMPONENT_RUN,
-    _COMPONENT_TYPE,
+    _COMPONENT_VISITS,
     ComponentPriority,
     PipelineBase,
 )
@@ -51,28 +49,8 @@ def _run_component(
         """
         instance: Component = component["instance"]
 
-        with tracing.tracer.trace(
-            _COMPONENT_RUN,
-            tags={
-                _COMPONENT_NAME: component_name,
-                _COMPONENT_TYPE: instance.__class__.__name__,
-                "haystack.component.input_types": {k: type(v).__name__ for k, v in inputs.items()},
-                "haystack.component.input_spec": {
-                    key: {
-                        "type": (value.type.__name__ if isinstance(value.type, type) else str(value.type)),
-                        "senders": value.senders,
-                    }
-                    for key, value in instance.__haystack_input__._sockets_dict.items()  # type: ignore
-                },
-                "haystack.component.output_spec": {
-                    key: {
-                        "type": (value.type.__name__ if isinstance(value.type, type) else str(value.type)),
-                        "receivers": value.receivers,
-                    }
-                    for key, value in instance.__haystack_output__._sockets_dict.items()  # type: ignore
-                },
-            },
-            parent_span=parent_span,
+        with PipelineBase._create_component_span(
+            component_name=component_name, instance=instance, inputs=inputs, parent_span=parent_span
         ) as span:
             # We deepcopy the inputs otherwise we might lose that information
             # when we delete them in case they're sent to other Components
@@ -87,7 +65,7 @@ def _run_component(
             if not isinstance(component_output, Mapping):
                 raise PipelineRuntimeError.from_invalid_output(component_name, instance.__class__, component_output)
 
-            span.set_tag("haystack.component.visits", component_visits[component_name])
+            span.set_tag(_COMPONENT_VISITS, component_visits[component_name])
             span.set_content_tag(_COMPONENT_OUTPUT, component_output)
 
             return cast(Dict[Any, Any], component_output)