From 7ed964230554569962a7b4f792e2f3a5c99f16f6 Mon Sep 17 00:00:00 2001 From: Jacques Verre Date: Sat, 23 Nov 2024 22:10:17 +0000 Subject: [PATCH 1/4] update_current_span now supports updating the output --- sdks/python/src/opik/api_objects/span.py | 18 +++++++++ sdks/python/src/opik/api_objects/trace.py | 18 +++++++++ .../opik/decorator/base_track_decorator.py | 2 +- .../unit/decorator/test_tracker_outputs.py | 39 +++++++++++++++++++ 4 files changed, 76 insertions(+), 1 deletion(-) diff --git a/sdks/python/src/opik/api_objects/span.py b/sdks/python/src/opik/api_objects/span.py index ec75b15012..99abcfd125 100644 --- a/sdks/python/src/opik/api_objects/span.py +++ b/sdks/python/src/opik/api_objects/span.py @@ -271,6 +271,12 @@ def update(self, **new_data: Any) -> "SpanData": if key == "metadata": self._update_metadata(value) continue + elif key == "output": + self._update_output(value) + continue + elif key == "input": + self._update_input(value) + continue self.__dict__[key] = value @@ -282,6 +288,18 @@ def _update_metadata(self, new_metadata: Dict[str, Any]) -> None: else: self.metadata = dict_utils.deepmerge(self.metadata, new_metadata) + def _update_output(self, new_output: Dict[str, Any]) -> None: + if self.output is None: + self.output = new_output + else: + self.output = dict_utils.deepmerge(self.output, new_output) + + def _update_input(self, new_input: Dict[str, Any]) -> None: + if self.input is None: + self.input = new_input + else: + self.input = dict_utils.deepmerge(self.input, new_input) + def init_end_time(self) -> "SpanData": self.end_time = datetime_helpers.local_timestamp() diff --git a/sdks/python/src/opik/api_objects/trace.py b/sdks/python/src/opik/api_objects/trace.py index 3a80c210c5..37bb8701df 100644 --- a/sdks/python/src/opik/api_objects/trace.py +++ b/sdks/python/src/opik/api_objects/trace.py @@ -242,6 +242,12 @@ def update(self, **new_data: Any) -> "TraceData": if key == "metadata": self._update_metadata(value) continue + elif key == "output": + self._update_output(value) + continue + elif key == "input": + self._update_input(value) + continue self.__dict__[key] = value @@ -253,6 +259,18 @@ def _update_metadata(self, new_metadata: Dict[str, Any]) -> None: else: self.metadata = dict_utils.deepmerge(self.metadata, new_metadata) + def _update_output(self, new_output: Dict[str, Any]) -> None: + if self.output is None: + self.output = new_output + else: + self.output = dict_utils.deepmerge(self.output, new_output) + + def _update_input(self, new_input: Dict[str, Any]) -> None: + if self.input is None: + self.input = new_input + else: + self.input = dict_utils.deepmerge(self.input, new_input) + def init_end_time(self) -> "TraceData": self.end_time = datetime_helpers.local_timestamp() return self diff --git a/sdks/python/src/opik/decorator/base_track_decorator.py b/sdks/python/src/opik/decorator/base_track_decorator.py index 814e240c2c..f1d551366d 100644 --- a/sdks/python/src/opik/decorator/base_track_decorator.py +++ b/sdks/python/src/opik/decorator/base_track_decorator.py @@ -399,7 +399,7 @@ def _after_call( ) client = opik_client.get_client_cached() - + print("span_data_to_end", span_data_to_end) span_data_to_end.init_end_time().update( **end_arguments.to_kwargs(), ) diff --git a/sdks/python/tests/unit/decorator/test_tracker_outputs.py b/sdks/python/tests/unit/decorator/test_tracker_outputs.py index aa542b1a62..3889fdb7c2 100644 --- a/sdks/python/tests/unit/decorator/test_tracker_outputs.py +++ b/sdks/python/tests/unit/decorator/test_tracker_outputs.py @@ -936,6 +936,44 @@ def f(x): assert_equal(EXPECTED_TRACE_TREE, fake_backend.trace_trees[0]) +def test_track__span_and_trace_output_updated_via_opik_context(fake_backend): + @tracker.track + def f(x): + opik_context.update_current_span( + output={"span-output-key": "span-output-value"}, + ) + opik_context.update_current_trace( + output={"trace-output-key": "trace-output-value"}, + ) + + return "f-output" + + f("f-input") + tracker.flush_tracker() + + EXPECTED_TRACE_TREE = TraceModel( + id=ANY_BUT_NONE, + name="f", + input={"x": "f-input"}, + output={"output": "f-output", "trace-output-key": "trace-output-value"}, + start_time=ANY_BUT_NONE, + end_time=ANY_BUT_NONE, + spans=[ + SpanModel( + id=ANY_BUT_NONE, + name="f", + input={"x": "f-input"}, + output={"output": "f-output", "span-output-key": "span-output-value"}, + start_time=ANY_BUT_NONE, + end_time=ANY_BUT_NONE, + spans=[], + ) + ], + ) + + assert len(fake_backend.trace_trees) == 1 + + assert_equal(EXPECTED_TRACE_TREE, fake_backend.trace_trees[0]) def test_track__span_and_trace_updated_via_opik_context_with_feedback_scores__feedback_scores_are_also_logged( fake_backend, @@ -989,6 +1027,7 @@ def f(x): assert_equal(EXPECTED_TRACE_TREE, fake_backend.trace_trees[0]) + def test_tracker__ignore_list_was_passed__ignored_inputs_are_not_logged(fake_backend): @tracker.track(ignore_arguments=["a", "c", "e", "unknown_argument"]) def f(a, b, c=3, d=4, e=5): From 8ef179656db7bf227c80717c4a8114a1ff7f065c Mon Sep 17 00:00:00 2001 From: Jacques Verre Date: Sat, 23 Nov 2024 22:18:12 +0000 Subject: [PATCH 2/4] Fix linter --- sdks/python/tests/unit/decorator/test_tracker_outputs.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/sdks/python/tests/unit/decorator/test_tracker_outputs.py b/sdks/python/tests/unit/decorator/test_tracker_outputs.py index 3889fdb7c2..30f58ea134 100644 --- a/sdks/python/tests/unit/decorator/test_tracker_outputs.py +++ b/sdks/python/tests/unit/decorator/test_tracker_outputs.py @@ -936,6 +936,7 @@ def f(x): assert_equal(EXPECTED_TRACE_TREE, fake_backend.trace_trees[0]) + def test_track__span_and_trace_output_updated_via_opik_context(fake_backend): @tracker.track def f(x): @@ -975,6 +976,7 @@ def f(x): assert_equal(EXPECTED_TRACE_TREE, fake_backend.trace_trees[0]) + def test_track__span_and_trace_updated_via_opik_context_with_feedback_scores__feedback_scores_are_also_logged( fake_backend, ): @@ -1027,7 +1029,6 @@ def f(x): assert_equal(EXPECTED_TRACE_TREE, fake_backend.trace_trees[0]) - def test_tracker__ignore_list_was_passed__ignored_inputs_are_not_logged(fake_backend): @tracker.track(ignore_arguments=["a", "c", "e", "unknown_argument"]) def f(a, b, c=3, d=4, e=5): From 8851945a3e4523a0272b5cee81a4b72e0252e217 Mon Sep 17 00:00:00 2001 From: Jacques Verre Date: Sun, 24 Nov 2024 13:29:54 +0000 Subject: [PATCH 3/4] Update end to end test --- sdks/python/tests/unit/decorator/test_tracker_outputs.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/sdks/python/tests/unit/decorator/test_tracker_outputs.py b/sdks/python/tests/unit/decorator/test_tracker_outputs.py index 30f58ea134..3c1f66761d 100644 --- a/sdks/python/tests/unit/decorator/test_tracker_outputs.py +++ b/sdks/python/tests/unit/decorator/test_tracker_outputs.py @@ -937,13 +937,15 @@ def f(x): assert_equal(EXPECTED_TRACE_TREE, fake_backend.trace_trees[0]) -def test_track__span_and_trace_output_updated_via_opik_context(fake_backend): +def test_track__span_and_trace_input_output_updated_via_opik_context(fake_backend): @tracker.track def f(x): opik_context.update_current_span( + input={"span-input-key": "span-input-value"}, output={"span-output-key": "span-output-value"}, ) opik_context.update_current_trace( + input={"trace-input-key": "trace-input-value"}, output={"trace-output-key": "trace-output-value"}, ) @@ -955,7 +957,7 @@ def f(x): EXPECTED_TRACE_TREE = TraceModel( id=ANY_BUT_NONE, name="f", - input={"x": "f-input"}, + input={"x": "f-input", "trace-input-key": "trace-input-value"}, output={"output": "f-output", "trace-output-key": "trace-output-value"}, start_time=ANY_BUT_NONE, end_time=ANY_BUT_NONE, @@ -963,7 +965,7 @@ def f(x): SpanModel( id=ANY_BUT_NONE, name="f", - input={"x": "f-input"}, + input={"x": "f-input", "span-input-key": "span-input-value"}, output={"output": "f-output", "span-output-key": "span-output-value"}, start_time=ANY_BUT_NONE, end_time=ANY_BUT_NONE, From 6228286bfbca5f2430c724ef0488d1efb1b03c51 Mon Sep 17 00:00:00 2001 From: Jacques Verre Date: Mon, 25 Nov 2024 10:30:01 +0000 Subject: [PATCH 4/4] Remove print statement --- sdks/python/src/opik/decorator/base_track_decorator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdks/python/src/opik/decorator/base_track_decorator.py b/sdks/python/src/opik/decorator/base_track_decorator.py index f1d551366d..814e240c2c 100644 --- a/sdks/python/src/opik/decorator/base_track_decorator.py +++ b/sdks/python/src/opik/decorator/base_track_decorator.py @@ -399,7 +399,7 @@ def _after_call( ) client = opik_client.get_client_cached() - print("span_data_to_end", span_data_to_end) + span_data_to_end.init_end_time().update( **end_arguments.to_kwargs(), )