fix(media): allow setting IO media via decorator update (#1005)

Author: hassiebp

langfuse/decorators/langfuse_decorator.py

@@ -446,13 +446,18 @@ def _handle_call_result(
             )
 
             end_time = observation_params["end_time"] or _get_timestamp()
-            raw_output = observation_params["output"] or (
-                result if result and capture_output else None
+
+            output = observation_params["output"] or (
+                # Serialize and deserialize to ensure proper JSON serialization.
+                # Objects are later serialized again so deserialization is necessary here to avoid unnecessary escaping of quotes.
+                json.loads(
+                    json.dumps(
+                        result if result and capture_output else None,
+                        cls=EventSerializer,
+                    )
+                )
             )
 
-            # Serialize and deserialize to ensure proper JSON serialization.
-            # Objects are later serialized again so deserialization is necessary here to avoid unnecessary escaping of quotes.
-            output = json.loads(json.dumps(raw_output, cls=EventSerializer))
             observation_params.update(end_time=end_time, output=output)
 
             if isinstance(observation, (StatefulSpanClient, StatefulGenerationClient)):

tests/test_decorators.py

@@ -1473,6 +1473,16 @@ def test_media():
     @observe()
     def main():
         langfuse_context.update_current_trace(
+            input={
+                "context": {
+                    "nested": media,
+                },
+            },
+            output={
+                "context": {
+                    "nested": media,
+                },
+            },
             metadata={
                 "context": {
                     "nested": media,
@@ -1486,6 +1496,14 @@ def main():
 
     trace_data = get_api().trace.get(mock_trace_id)
 
+    assert (
+        "@@@langfuseMedia:type=application/pdf|id="
+        in trace_data.input["context"]["nested"]
+    )
+    assert (
+        "@@@langfuseMedia:type=application/pdf|id="
+        in trace_data.output["context"]["nested"]
+    )
     assert (
         "@@@langfuseMedia:type=application/pdf|id="
         in trace_data.metadata["context"]["nested"]