update_current_span now supports updating the output

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()
 

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

sdks/python/tests/unit/decorator/test_tracker_outputs.py

@@ -937,6 +937,48 @@ def f(x):
     assert_equal(EXPECTED_TRACE_TREE, fake_backend.trace_trees[0])
 
 
+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"},
+        )
+
+        return "f-output"
+
+    f("f-input")
+    tracker.flush_tracker()
+
+    EXPECTED_TRACE_TREE = TraceModel(
+        id=ANY_BUT_NONE,
+        name="f",
+        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,
+        spans=[
+            SpanModel(
+                id=ANY_BUT_NONE,
+                name="f",
+                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,
+                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,
 ):