-
Notifications
You must be signed in to change notification settings - Fork 2.3k
Python: Capture workflow telemetry input and output #7565
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -19,7 +19,12 @@ | |
| from .._sessions import ContextProvider | ||
| from .._types import ResponseStream | ||
| from ..exceptions import WorkflowException | ||
| from ..observability import OtelAttr, capture_exception, create_workflow_span | ||
| from ..observability import ( | ||
| OtelAttr, | ||
| _set_sensitive_span_attributes, # pyright: ignore[reportPrivateUsage] | ||
| capture_exception, | ||
| create_workflow_span, | ||
| ) | ||
| from ._checkpoint import CheckpointStorage | ||
| from ._const import DEFAULT_MAX_ITERATIONS, GLOBAL_KWARGS_KEY, INTERNAL_SOURCE_ID, WORKFLOW_RUN_KWARGS_KEY | ||
| from ._edge import ( | ||
|
|
@@ -478,6 +483,7 @@ def get_executors_list(self) -> list[Executor]: | |
| async def _run_workflow_with_tracing( | ||
| self, | ||
| initial_executor_fn: Callable[[], Awaitable[None]] | None = None, | ||
| telemetry_input: Any | None = None, | ||
| is_continuation: bool = False, | ||
| streaming: bool = False, | ||
| function_invocation_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, | ||
|
|
@@ -490,6 +496,7 @@ async def _run_workflow_with_tracing( | |
|
|
||
| Args: | ||
| initial_executor_fn: Optional function to execute initial executor. | ||
| telemetry_input: Input payload for this run, captured only when sensitive telemetry is enabled. | ||
| is_continuation: True when this run is a continuation of prior | ||
| work (a checkpoint restore or a responses-only replay) rather | ||
| than a fresh new turn delivered via the start executor with | ||
|
|
@@ -517,9 +524,20 @@ async def _run_workflow_with_tracing( | |
| OtelAttr.WORKFLOW_RUN_SPAN, | ||
| attributes, | ||
| ) as span: | ||
| from ..observability import OBSERVABILITY_SETTINGS | ||
|
|
||
| capture_workflow_io = OBSERVABILITY_SETTINGS.SENSITIVE_DATA_ENABLED and span.is_recording() | ||
| saw_request = False | ||
| emitted_in_progress_pending = False | ||
| workflow_outputs: list[Any] | None = [] if capture_workflow_io else None | ||
| try: | ||
| if capture_workflow_io and telemetry_input is not None: | ||
| _set_sensitive_span_attributes( | ||
| span, | ||
| telemetry_input, | ||
| (OtelAttr.INPUT_VALUE,), | ||
| (OtelAttr.INPUT_MIME_TYPE,), | ||
| ) | ||
| # Add workflow started event (telemetry + surface state to consumers) | ||
| span.add_event(OtelAttr.WORKFLOW_STARTED) | ||
| # Emit explicit start/status events to the stream | ||
|
|
@@ -577,6 +595,8 @@ async def _run_workflow_with_tracing( | |
| # Track request events for final status determination | ||
| if event.type == "request_info": | ||
| saw_request = True | ||
| elif workflow_outputs is not None and event.type == "output": | ||
| workflow_outputs.append(event.data) | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Could we please avoid retaining and serializing every streamed output as one span attribute? With sensitive telemetry enabled, |
||
| yield event | ||
|
|
||
| if event.type == "request_info" and not emitted_in_progress_pending: | ||
|
|
@@ -597,6 +617,13 @@ async def _run_workflow_with_tracing( | |
| terminal_status = WorkflowEvent.status(self._status) | ||
| yield terminal_status | ||
|
|
||
| if workflow_outputs: | ||
| _set_sensitive_span_attributes( | ||
| span, | ||
| workflow_outputs, | ||
| (OtelAttr.OUTPUT_VALUE,), | ||
| (OtelAttr.OUTPUT_MIME_TYPE,), | ||
| ) | ||
| span.add_event(OtelAttr.WORKFLOW_COMPLETED) | ||
| except Exception as exc: | ||
| # Drain any pending events (for example, executor_failed) before yielding failed event | ||
|
|
@@ -872,9 +899,13 @@ async def _run_core( | |
| ) | ||
|
|
||
| initial_executor_fn = self._resolve_execution_mode(message, responses, checkpoint_id, checkpoint_storage) | ||
| telemetry_input = message.data if isinstance(message, WorkflowMessage) else message | ||
| if telemetry_input is None: | ||
| telemetry_input = responses | ||
|
|
||
| async for event in self._run_workflow_with_tracing( | ||
| initial_executor_fn=initial_executor_fn, | ||
| telemetry_input=telemetry_input, | ||
| is_continuation=(message is None), | ||
| streaming=streaming, | ||
| function_invocation_kwargs=function_invocation_kwargs, | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -44,6 +44,7 @@ | |
| from . import __version__ as version_info | ||
| from ._serialization import ( | ||
| _is_serialization_protocol, # pyright: ignore[reportPrivateUsage] | ||
| make_json_safe, | ||
| ) | ||
| from ._settings import load_settings | ||
|
|
||
|
|
@@ -304,6 +305,8 @@ class OtelAttr(str, Enum): | |
| EXECUTOR_PROCESS_SPAN = "executor.process" | ||
| EXECUTOR_ID = "executor.id" | ||
| EXECUTOR_TYPE = "executor.type" | ||
| EXECUTOR_INPUT = "executor.input" | ||
| EXECUTOR_OUTPUT = "executor.output" | ||
| # Edge group attributes | ||
| EDGE_GROUP_PROCESS_SPAN = "edge_group.process" | ||
| EDGE_GROUP_TYPE = "edge_group.type" | ||
|
|
@@ -317,6 +320,15 @@ class OtelAttr(str, Enum): | |
| MESSAGE_TYPE = "message.type" | ||
| MESSAGE_PAYLOAD_TYPE = "message.payload_type" | ||
| MESSAGE_DESTINATION_EXECUTOR_ID = "message.destination_executor_id" | ||
| MESSAGE_CONTENT = "message.content" | ||
|
|
||
| # OpenInference attributes for vendor-neutral input/output ingestion. | ||
| # https://arize-ai.github.io/openinference/spec/semantic_conventions.html | ||
| INPUT_VALUE = "input.value" | ||
| INPUT_MIME_TYPE = "input.mime_type" | ||
| OUTPUT_VALUE = "output.value" | ||
| OUTPUT_MIME_TYPE = "output.mime_type" | ||
| JSON_MIME_TYPE = "application/json" | ||
|
|
||
| # Activity events | ||
| EVENT_NAME = "event.name" | ||
|
|
@@ -356,6 +368,32 @@ def __str__(self) -> str: | |
| return self.value | ||
|
|
||
|
|
||
| def _serialize_for_telemetry(value: Any) -> str: | ||
| """Serialize heterogeneous telemetry payloads without affecting application execution.""" | ||
| try: | ||
| return json.dumps(make_json_safe(value), ensure_ascii=False, allow_nan=False) | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
|
||
| except Exception: | ||
| value_type = f"{type(value).__module__}.{type(value).__qualname__}" | ||
| return json.dumps(f"[Unserializable: {value_type}]", ensure_ascii=False) | ||
|
|
||
|
|
||
| def _set_sensitive_span_attributes( # pyright: ignore[reportUnusedFunction] | ||
| span: trace.Span, | ||
| value: Any, | ||
| value_attributes: Sequence[str | OtelAttr], | ||
| mime_type_attributes: Sequence[str | OtelAttr] = (), | ||
| ) -> None: | ||
| """Set serialized payload attributes only when sensitive telemetry is enabled.""" | ||
| if not OBSERVABILITY_SETTINGS.SENSITIVE_DATA_ENABLED or not span.is_recording(): | ||
| return | ||
|
|
||
| serialized_value = _serialize_for_telemetry(value) | ||
| for attribute in value_attributes: | ||
| span.set_attribute(attribute, serialized_value) | ||
| for attribute in mime_type_attributes: | ||
| span.set_attribute(attribute, OtelAttr.JSON_MIME_TYPE) | ||
|
|
||
|
|
||
| ROLE_EVENT_MAP = { | ||
| "system": OtelAttr.SYSTEM_MESSAGE, | ||
| "user": OtelAttr.USER_MESSAGE, | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Output capture only runs in the successful convergence loop and the attribute is written after that loop completes. If an executor queues an output and then fails, the exception path drains and yields that output without recording it; similarly, closing the stream after an output skips finalization. This leaves
output.valueabsent even though the caller received workflow output. Please capture drained output events as well and finalize the accumulated attribute on failure or generator close without changing the workflow's output-designation rules.