Skip to content

Commit

Permalink
wip: composite_tracer
Browse files Browse the repository at this point in the history
  • Loading branch information
NiklasKoehneckeAA committed May 16, 2024
1 parent 6e3ca61 commit 7e4bf30
Show file tree
Hide file tree
Showing 2 changed files with 9 additions and 7 deletions.
14 changes: 8 additions & 6 deletions src/intelligence_layer/core/tracer/composite_tracer.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
from typing import Generic, Optional, Sequence, TypeVar

from intelligence_layer.core.tracer.tracer import (
ExportedSpan,
PydanticSerializable,
Span,
TaskSpan,
Expand Down Expand Up @@ -41,30 +42,31 @@ 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(
self,
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.
Expand Down
2 changes: 1 addition & 1 deletion tests/core/tracer/test_composite_tracer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down

0 comments on commit 7e4bf30

Please sign in to comment.