From 902bcfb579e7300eb6fc79b0d0945e7f7fc64256 Mon Sep 17 00:00:00 2001 From: Gustavo Cid Date: Tue, 22 Jul 2025 15:38:07 -0300 Subject: [PATCH] feat: allow output overwrite --- src/openlayer/lib/tracing/tracer.py | 24 +++++++++++++++++++++++- 1 file changed, 23 insertions(+), 1 deletion(-) diff --git a/src/openlayer/lib/tracing/tracer.py b/src/openlayer/lib/tracing/tracer.py index 83af81fb..0749fe90 100644 --- a/src/openlayer/lib/tracing/tracer.py +++ b/src/openlayer/lib/tracing/tracer.py @@ -371,6 +371,23 @@ def sync_wrapper(*func_args, **func_kwargs): return decorator +def log_output(output: Any) -> None: + """Logs output information to the current step of the trace. + + This will overwrite the output of the currently active step instead of + relying on the returned object from the traced function. + + Args: + output: The output value to log to the current step. + """ + current_step = get_current_step() + if current_step: + logger.debug("Logging output to current step: %s", output) + current_step.log(output=output, metadata={"manual_output_logged": True}) + else: + logger.warning("No current step found to log output.") + + def log_context(context: List[str]) -> None: """Logs context information to the current step of the trace. @@ -562,9 +579,14 @@ def _finalize_step_logging( if step.latency is None: step.latency = (step.end_time - start_time) * 1000 # in ms + # Check if manual output was logged + if step.metadata.get("manual_output_logged"): + logger.debug("Using manually logged output for step: %s", step.name) + else: + step.log(output=output) + step.log( inputs=inputs, - output=output, end_time=step.end_time, latency=step.latency, )