From 0c7575da64549a0642a6abbae73db6bef12fd6f4 Mon Sep 17 00:00:00 2001 From: Florian Schepers Date: Wed, 22 May 2024 12:05:42 +0200 Subject: [PATCH] Adapt File Tracer --- .../core/tracer/composite_tracer.py | 7 +++- .../core/tracer/file_tracer.py | 23 ++---------- .../core/tracer/in_memory_tracer.py | 14 ++++--- .../core/tracer/persistent_tracer.py | 37 +++++++++---------- src/intelligence_layer/core/tracer/tracer.py | 2 +- .../fixtures/old_file_trace_format.jsonl | 11 ++++++ tests/core/tracer/test_composite_tracer.py | 6 +-- tests/core/tracer/test_file_tracer.py | 27 ++++++++------ tests/core/tracer/test_in_memory_tracer.py | 1 + tests/core/tracer/test_tracer.py | 4 +- 10 files changed, 68 insertions(+), 64 deletions(-) create mode 100644 tests/core/tracer/fixtures/old_file_trace_format.jsonl diff --git a/src/intelligence_layer/core/tracer/composite_tracer.py b/src/intelligence_layer/core/tracer/composite_tracer.py index a8a255040..5dd9f1464 100644 --- a/src/intelligence_layer/core/tracer/composite_tracer.py +++ b/src/intelligence_layer/core/tracer/composite_tracer.py @@ -2,13 +2,13 @@ from typing import Generic, Optional, Sequence, TypeVar from intelligence_layer.core.tracer.tracer import ( + Context, ExportedSpan, PydanticSerializable, Span, TaskSpan, Tracer, utc_now, - Context ) TracerVar = TypeVar("TracerVar", bound=Tracer) @@ -73,7 +73,10 @@ class CompositeSpan(Generic[SpanVar], CompositeTracer[SpanVar], Span): Args: tracers: spans that will be forwarded all subsequent log and span calls. """ - def __init__(self, tracers: Sequence[SpanVar],context: Optional[Context] = None) -> None: + + def __init__( + self, tracers: Sequence[SpanVar], context: Optional[Context] = None + ) -> None: CompositeTracer.__init__(self, tracers) Span.__init__(self, context=context) diff --git a/src/intelligence_layer/core/tracer/file_tracer.py b/src/intelligence_layer/core/tracer/file_tracer.py index e2c9209cc..4615e9227 100644 --- a/src/intelligence_layer/core/tracer/file_tracer.py +++ b/src/intelligence_layer/core/tracer/file_tracer.py @@ -1,7 +1,7 @@ from datetime import datetime from json import loads from pathlib import Path -from typing import Optional, Sequence +from typing import Optional from pydantic import BaseModel @@ -11,22 +11,8 @@ PersistentTaskSpan, PersistentTracer, ) -from intelligence_layer.core.tracer.tracer import LogLine, PydanticSerializable -from intelligence_layer.core.tracer.tracer import ( - Context, - Event, - ExportedSpan, - ExportedSpanList, - LogEntry, - PydanticSerializable, - Span, - SpanAttributes, - TaskSpan, - TaskSpanAttributes, - Tracer, - _render_log_value, - utc_now, -) +from intelligence_layer.core.tracer.tracer import Context, LogLine, PydanticSerializable + class FileTracer(PersistentTracer): """A `Tracer` that logs to a file. @@ -91,16 +77,15 @@ def trace(self, trace_id: Optional[str] = None) -> InMemoryTracer: return self._parse_log(filtered_traces) - class FileSpan(PersistentSpan, FileTracer): """A `Span` created by `FileTracer.span`.""" def __init__(self, log_file_path: Path, context: Optional[Context] = None) -> None: PersistentSpan.__init__(self, context=context) FileTracer.__init__(self, log_file_path=log_file_path) - class FileTaskSpan(PersistentTaskSpan, FileSpan): """A `TaskSpan` created by `FileTracer.task_span`.""" + pass diff --git a/src/intelligence_layer/core/tracer/in_memory_tracer.py b/src/intelligence_layer/core/tracer/in_memory_tracer.py index 124cef20b..0923dc440 100644 --- a/src/intelligence_layer/core/tracer/in_memory_tracer.py +++ b/src/intelligence_layer/core/tracer/in_memory_tracer.py @@ -5,7 +5,7 @@ import requests import rich -from pydantic import BaseModel, Field, SerializeAsAny +from pydantic import SerializeAsAny from requests import HTTPError from rich.tree import Tree @@ -246,9 +246,11 @@ def start_task(self, log_line: LogLine) -> None: name=start_task.name, input=start_task.input, start_timestamp=start_task.start, - context=Context( + context=Context( trace_id=start_task.trace_id, span_id=str(start_task.parent) - ) if start_task.trace_id != str(start_task.uuid) else None + ) + if start_task.trace_id != str(start_task.uuid) + else None, ) # if root, also change the trace id if converted_span.context.trace_id == converted_span.context.span_id: @@ -271,13 +273,15 @@ def start_span(self, log_line: LogLine) -> None: start_timestamp=start_span.start, context=Context( trace_id=start_span.trace_id, span_id=str(start_span.parent) - ) if start_span.trace_id != str(start_span.uuid) else None + ) + if start_span.trace_id != str(start_span.uuid) + else None, ) # if root, also change the trace id if converted_span.context.trace_id == converted_span.context.span_id: converted_span.context.trace_id = str(start_span.uuid) converted_span.context.span_id = str(start_span.uuid) - + self.tracers.get(start_span.parent, self.root).entries.append(converted_span) self.tracers[start_span.uuid] = converted_span self.spans[start_span.uuid] = converted_span diff --git a/src/intelligence_layer/core/tracer/persistent_tracer.py b/src/intelligence_layer/core/tracer/persistent_tracer.py index 0cd5435f1..453e35ca5 100644 --- a/src/intelligence_layer/core/tracer/persistent_tracer.py +++ b/src/intelligence_layer/core/tracer/persistent_tracer.py @@ -9,6 +9,7 @@ from intelligence_layer.core.tracer.tracer import ( EndSpan, EndTask, + ExportedSpan, LogLine, PlainEntry, PydanticSerializable, @@ -20,26 +21,11 @@ utc_now, ) -from intelligence_layer.core.tracer.tracer import ( - Context, - Event, - ExportedSpan, - ExportedSpanList, - LogEntry, - PydanticSerializable, - Span, - SpanAttributes, - TaskSpan, - TaskSpanAttributes, - Tracer, - _render_log_value, - utc_now, -) class PersistentTracer(Tracer, ABC): def __init__(self) -> None: self.current_id = uuid4() - + @abstractmethod def _log_entry(self, id: str, entry: BaseModel) -> None: pass @@ -77,7 +63,9 @@ def _log_task( task_span.context.trace_id, StartTask( uuid=task_span.context.span_id, - parent=self.context.span_id if self.context else task_span.context.trace_id, + parent=self.context.span_id + if self.context + else task_span.context.trace_id, name=task_name, start=timestamp or utc_now(), input=input, @@ -89,7 +77,9 @@ def _log_task( task_span.context.trace_id, StartTask( uuid=task_span.context.span_id, - parent=self.context.span_id if self.context else task_span.context.trace_id, + parent=self.context.span_id + if self.context + else task_span.context.trace_id, name=task_name, start=timestamp or utc_now(), input=error.description, @@ -151,7 +141,10 @@ def log( def end(self, timestamp: Optional[datetime] = None) -> None: if not self.end_timestamp: self.end_timestamp = timestamp or utc_now() - self._log_entry(self.context.trace_id, EndSpan(uuid=self.context.span_id, end=self.end_timestamp)) + self._log_entry( + self.context.trace_id, + EndSpan(uuid=self.context.span_id, end=self.end_timestamp), + ) class PersistentTaskSpan(TaskSpan, PersistentSpan, ABC): @@ -165,7 +158,11 @@ def end(self, timestamp: Optional[datetime] = None) -> None: self.end_timestamp = timestamp or utc_now() self._log_entry( self.context.trace_id, - EndTask(uuid=self.context.span_id, end=self.end_timestamp, output=self.output), + EndTask( + uuid=self.context.span_id, + end=self.end_timestamp, + output=self.output, + ), ) diff --git a/src/intelligence_layer/core/tracer/tracer.py b/src/intelligence_layer/core/tracer/tracer.py index d726540c2..9102ae704 100644 --- a/src/intelligence_layer/core/tracer/tracer.py +++ b/src/intelligence_layer/core/tracer/tracer.py @@ -204,7 +204,7 @@ class Span(Tracer, AbstractContextManager["Span"]): """ def __init__(self, context: Optional[Context] = None): - #super().__init__() + # super().__init__() span_id = str(uuid4()) if context is None: trace_id = span_id diff --git a/tests/core/tracer/fixtures/old_file_trace_format.jsonl b/tests/core/tracer/fixtures/old_file_trace_format.jsonl new file mode 100644 index 000000000..bf1fb390e --- /dev/null +++ b/tests/core/tracer/fixtures/old_file_trace_format.jsonl @@ -0,0 +1,11 @@ +{"trace_id":"00b96d92-d1d1-454a-b8e9-f67e8541759a","entry_type":"StartTask","entry":{"uuid":"41528209-1b78-4785-a00d-7f65af1bb09c","parent":"75e79a11-1a26-4731-8b49-ef8634c352ed","name":"TestTask","start":"2024-05-22T09:43:37.428758Z","input":"input","trace_id":"00b96d92-d1d1-454a-b8e9-f67e8541759a"}} +{"trace_id":"00b96d92-d1d1-454a-b8e9-f67e8541759a","entry_type":"StartSpan","entry":{"uuid":"ad1ed79b-6ad6-4ea5-8ee8-26be4055e228","parent":"41528209-1b78-4785-a00d-7f65af1bb09c","name":"span","start":"2024-05-22T09:43:37.429448Z","trace_id":"00b96d92-d1d1-454a-b8e9-f67e8541759a"}} +{"trace_id":"00b96d92-d1d1-454a-b8e9-f67e8541759a","entry_type":"PlainEntry","entry":{"message":"message","value":"a value","timestamp":"2024-05-22T09:43:37.429503Z","parent":"ad1ed79b-6ad6-4ea5-8ee8-26be4055e228","trace_id":"00b96d92-d1d1-454a-b8e9-f67e8541759a"}} +{"trace_id":"00b96d92-d1d1-454a-b8e9-f67e8541759a","entry_type":"StartTask","entry":{"uuid":"e8cca541-57a8-440a-b848-7c3b33a97f52","parent":"ad1ed79b-6ad6-4ea5-8ee8-26be4055e228","name":"TestSubTask","start":"2024-05-22T09:43:37.429561Z","input":null,"trace_id":"00b96d92-d1d1-454a-b8e9-f67e8541759a"}} +{"trace_id":"00b96d92-d1d1-454a-b8e9-f67e8541759a","entry_type":"PlainEntry","entry":{"message":"subtask","value":"value","timestamp":"2024-05-22T09:43:37.429605Z","parent":"e8cca541-57a8-440a-b848-7c3b33a97f52","trace_id":"00b96d92-d1d1-454a-b8e9-f67e8541759a"}} +{"trace_id":"00b96d92-d1d1-454a-b8e9-f67e8541759a","entry_type":"EndTask","entry":{"uuid":"e8cca541-57a8-440a-b848-7c3b33a97f52","end":"2024-05-22T09:43:37.429647Z","output":null}} +{"trace_id":"00b96d92-d1d1-454a-b8e9-f67e8541759a","entry_type":"EndSpan","entry":{"uuid":"ad1ed79b-6ad6-4ea5-8ee8-26be4055e228","end":"2024-05-22T09:43:37.429687Z"}} +{"trace_id":"00b96d92-d1d1-454a-b8e9-f67e8541759a","entry_type":"StartTask","entry":{"uuid":"8840185c-2019-4105-9178-1b0e20ab6388","parent":"41528209-1b78-4785-a00d-7f65af1bb09c","name":"TestSubTask","start":"2024-05-22T09:43:37.429728Z","input":null,"trace_id":"00b96d92-d1d1-454a-b8e9-f67e8541759a"}} +{"trace_id":"00b96d92-d1d1-454a-b8e9-f67e8541759a","entry_type":"PlainEntry","entry":{"message":"subtask","value":"value","timestamp":"2024-05-22T09:43:37.429768Z","parent":"8840185c-2019-4105-9178-1b0e20ab6388","trace_id":"00b96d92-d1d1-454a-b8e9-f67e8541759a"}} +{"trace_id":"00b96d92-d1d1-454a-b8e9-f67e8541759a","entry_type":"EndTask","entry":{"uuid":"8840185c-2019-4105-9178-1b0e20ab6388","end":"2024-05-22T09:43:37.429806Z","output":null}} +{"trace_id":"00b96d92-d1d1-454a-b8e9-f67e8541759a","entry_type":"EndTask","entry":{"uuid":"41528209-1b78-4785-a00d-7f65af1bb09c","end":"2024-05-22T09:43:37.429842Z","output":"output"}} diff --git a/tests/core/tracer/test_composite_tracer.py b/tests/core/tracer/test_composite_tracer.py index 06cc3bae2..17c5c0613 100644 --- a/tests/core/tracer/test_composite_tracer.py +++ b/tests/core/tracer/test_composite_tracer.py @@ -1,8 +1,4 @@ -from intelligence_layer.core import ( - CompositeTracer, - InMemoryTracer, - Task, -) +from intelligence_layer.core import CompositeTracer, InMemoryTracer, Task def test_composite_tracer(test_task: Task[str, str]) -> None: diff --git a/tests/core/tracer/test_file_tracer.py b/tests/core/tracer/test_file_tracer.py index afc55ef5b..3cee09ff2 100644 --- a/tests/core/tracer/test_file_tracer.py +++ b/tests/core/tracer/test_file_tracer.py @@ -5,6 +5,7 @@ from pytest import fixture from intelligence_layer.core import CompositeTracer, FileTracer, InMemoryTracer, Task +from intelligence_layer.core.tracer.in_memory_tracer import InMemoryTaskSpan from intelligence_layer.core.tracer.persistent_tracer import TracerLogEntryFailed @@ -13,17 +14,6 @@ def file_tracer(tmp_path: Path) -> FileTracer: return FileTracer(tmp_path / "log.log") -def test_file_tracer(file_tracer: FileTracer, test_task: Task[str, str]) -> None: - input = "input" - expected = InMemoryTracer() - - test_task.run(input, file_tracer) - - log_tree = file_tracer.trace() - trace_1 = log_tree.export_for_viewing() - - - def test_file_tracer_retrieves_correct_trace( file_tracer: FileTracer, test_task: Task[str, str] ) -> None: @@ -59,3 +49,18 @@ def test_file_tracer_raises_non_log_entry_failed_exceptions( file_tracer.task_span( task_name="mock_task_name", input="42", timestamp=None, trace_id="21" ) + + +def test_file_tracer_is_backwards_compatible() -> None: + current_file_location = Path(__file__) + file_tracer = FileTracer( + current_file_location.parent / "fixtures/old_file_trace_format.jsonl" + ) + tracer = file_tracer.trace() + + assert len(tracer.entries) == 1 + task_span = tracer.entries[0] + assert isinstance(task_span, InMemoryTaskSpan) + assert task_span.input == "input" + assert task_span.start_timestamp and task_span.end_timestamp + assert task_span.start_timestamp < task_span.end_timestamp diff --git a/tests/core/tracer/test_in_memory_tracer.py b/tests/core/tracer/test_in_memory_tracer.py index afebd2760..bc9c16695 100644 --- a/tests/core/tracer/test_in_memory_tracer.py +++ b/tests/core/tracer/test_in_memory_tracer.py @@ -106,6 +106,7 @@ def test_task_span_records_error_value() -> None: def test_task_automatically_logs_input_and_output( test_task: Task[str, str], ) -> None: + input = "input" tracer = InMemoryTracer() output = test_task.run(input=input, tracer=tracer) diff --git a/tests/core/tracer/test_tracer.py b/tests/core/tracer/test_tracer.py index 749857c08..b56acc442 100644 --- a/tests/core/tracer/test_tracer.py +++ b/tests/core/tracer/test_tracer.py @@ -50,7 +50,9 @@ def test_tracer_exports_spans_to_unified_format( assert len(span.events) == 1 log = span.events[0] assert log.message == "test" - assert log.body == dummy_object or DummyObject.model_validate(log.body) == dummy_object + assert ( + log.body == dummy_object or DummyObject.model_validate(log.body) == dummy_object + ) assert span.start_time < log.timestamp < span.end_time