diff --git a/src/intelligence_layer/core/tracer/composite_tracer.py b/src/intelligence_layer/core/tracer/composite_tracer.py index 609800a36..f9c581c26 100644 --- a/src/intelligence_layer/core/tracer/composite_tracer.py +++ b/src/intelligence_layer/core/tracer/composite_tracer.py @@ -2,6 +2,7 @@ from typing import Generic, Optional, Sequence, TypeVar from intelligence_layer.core.tracer.tracer import ( + ExportedSpan, PydanticSerializable, Span, TaskSpan, @@ -41,12 +42,10 @@ def span( self, name: str, timestamp: Optional[datetime] = None, - trace_id: Optional[str] = None, ) -> "CompositeSpan[Span]": timestamp = timestamp or utc_now() - trace_id = self.ensure_id(trace_id) return CompositeSpan( - [tracer.span(name, timestamp, trace_id) for tracer in self.tracers] + [tracer.span(name, timestamp) for tracer in self.tracers] ) def task_span( @@ -54,17 +53,20 @@ def task_span( task_name: str, input: PydanticSerializable, timestamp: Optional[datetime] = None, - trace_id: Optional[str] = None, ) -> "CompositeTaskSpan": timestamp = timestamp or utc_now() - trace_id = self.ensure_id(trace_id) return CompositeTaskSpan( [ - tracer.task_span(task_name, input, timestamp, trace_id) + tracer.task_span(task_name, input, timestamp) for tracer in self.tracers ] ) + def export_for_viewing(self) -> Sequence[ExportedSpan]: + if len(self.tracers) > 0: + return self.tracers[0].export_for_viewing() + return [] + class CompositeSpan(Generic[SpanVar], CompositeTracer[SpanVar], Span): """A :class:`Span` that allows for recording to multiple spans simultaneously. diff --git a/tests/core/tracer/test_composite_tracer.py b/tests/core/tracer/test_composite_tracer.py index 4a33de925..1706c1742 100644 --- a/tests/core/tracer/test_composite_tracer.py +++ b/tests/core/tracer/test_composite_tracer.py @@ -12,7 +12,7 @@ def test_composite_tracer(test_task: Task[str, str]) -> None: tracer2 = InMemoryTracer() test_task.run(input=input, tracer=CompositeTracer([tracer1, tracer2])) - assert tracer1 == tracer2 + assert tracer1.export_for_viewing() == tracer2.export_for_viewing() def test_composite_tracer_id_consistent_across_children(