From 379e3b9a0094be4d5ea26aa69bafa72d9daa7173 Mon Sep 17 00:00:00 2001 From: Eric Zhu Date: Tue, 26 Aug 2025 22:09:16 -0700 Subject: [PATCH] Python: Tracing for workflows (#480) * workflow tracing design doc * add tracing implementation for workflow * fix bug caused by double wrapping of sub workflow request * add unit tests for tracing * add documentation for workflow tracing * remove unnecessary file * update aspire command * fix tests * proper serialization of subworkflows and add workflow.definition * add serialization test * fix subworkflow serialization * workflow_id --> id * update workflow sample to address comments * update naming; use costant * use NoOpTracer instead of nullcontext * use span event instead of attribtutes for status * fix typing * add workflow.build span * rename methods for clarity * ensure all source trace contexts are propagated in fan in --- .../agent_framework_workflow/__init__.py | 8 + .../agent_framework_workflow/_edge_runner.py | 50 +- .../agent_framework_workflow/_executor.py | 90 +-- .../agent_framework_workflow/_runner.py | 19 +- .../_runner_context.py | 27 +- .../agent_framework_workflow/_telemetry.py | 213 +++++++ .../agent_framework_workflow/_workflow.py | 268 ++++++--- .../_workflow_context.py | 32 +- .../workflow/tests/test_serialization.py | 109 +++- .../packages/workflow/tests/test_tracing.py | 532 ++++++++++++++++++ .../packages/workflow/tests/test_workflow.py | 2 +- .../getting_started/telemetry/.env.example | 1 + .../getting_started/telemetry/README.md | 69 ++- .../getting_started/telemetry/workflow.py | 253 +++++++++ 14 files changed, 1533 insertions(+), 140 deletions(-) create mode 100644 python/packages/workflow/agent_framework_workflow/_telemetry.py create mode 100644 python/packages/workflow/tests/test_tracing.py create mode 100644 python/samples/getting_started/telemetry/workflow.py diff --git a/python/packages/workflow/agent_framework_workflow/__init__.py b/python/packages/workflow/agent_framework_workflow/__init__.py index aace09fc67..12b8ce6e4c 100644 --- a/python/packages/workflow/agent_framework_workflow/__init__.py +++ b/python/packages/workflow/agent_framework_workflow/__init__.py @@ -105,3 +105,11 @@ __all__ = [ "intercepts_request", "validate_workflow_graph", ] + + +# Rebuild models to resolve forward references after all imports are complete +import contextlib + +with contextlib.suppress(AttributeError, TypeError, ValueError): + # Rebuild WorkflowExecutor to resolve Workflow forward reference + WorkflowExecutor.model_rebuild() diff --git a/python/packages/workflow/agent_framework_workflow/_edge_runner.py b/python/packages/workflow/agent_framework_workflow/_edge_runner.py index 184cf0b940..7ccc50199d 100644 --- a/python/packages/workflow/agent_framework_workflow/_edge_runner.py +++ b/python/packages/workflow/agent_framework_workflow/_edge_runner.py @@ -49,14 +49,36 @@ class EdgeRunner(ABC): return self._executors[executor_id].can_handle(message_data) async def _execute_on_target( - self, target_id: str, source_id: str, message_data: Any, shared_state: SharedState, ctx: RunnerContext + self, target_id: str, source_id: str, message: Message, shared_state: SharedState, ctx: RunnerContext ) -> None: - """Execute a message on a target executor.""" + """Execute a message on a target executor with trace context.""" if target_id not in self._executors: raise RuntimeError(f"Target executor {target_id} not found.") target_executor = self._executors[target_id] - await target_executor.execute(message_data, WorkflowContext(target_id, [source_id], shared_state, ctx)) + + # Handle both old single trace context format and new multiple trace contexts format + trace_contexts = getattr(message, "trace_contexts", None) + source_span_ids = getattr(message, "source_span_ids", None) + + # Backwards compatibility: if old format is used, convert to new format + if trace_contexts is None and hasattr(message, "trace_context") and message.trace_context: + trace_contexts = [message.trace_context] + if source_span_ids is None and hasattr(message, "source_span_id") and message.source_span_id: + source_span_ids = [message.source_span_id] + + # Create WorkflowContext with trace contexts from message + workflow_context: WorkflowContext[Any] = WorkflowContext( + target_id, + [source_id], + shared_state, + ctx, + trace_contexts=trace_contexts, # Pass trace contexts to WorkflowContext + source_span_ids=source_span_ids, # Pass source span IDs for linking + ) + + # Execute with trace context in WorkflowContext + await target_executor.execute(message.data, workflow_context) class SingleEdgeRunner(EdgeRunner): @@ -73,9 +95,7 @@ class SingleEdgeRunner(EdgeRunner): if self._can_handle(self._edge.target_id, message.data): if self._edge.should_route(message.data): - await self._execute_on_target( - self._edge.target_id, self._edge.source_id, message.data, shared_state, ctx - ) + await self._execute_on_target(self._edge.target_id, self._edge.source_id, message, shared_state, ctx) return True return False @@ -108,7 +128,7 @@ class FanOutEdgeRunner(EdgeRunner): edge = self._target_map.get(message.target_id) if edge and self._can_handle(edge.target_id, message.data): if edge.should_route(message.data): - await self._execute_on_target(edge.target_id, edge.source_id, message.data, shared_state, ctx) + await self._execute_on_target(edge.target_id, edge.source_id, message, shared_state, ctx) return True return False @@ -117,7 +137,7 @@ class FanOutEdgeRunner(EdgeRunner): """Send the message to the edge.""" if self._can_handle(edge.target_id, message.data): if edge.should_route(message.data): - await self._execute_on_target(edge.target_id, edge.source_id, message.data, shared_state, ctx) + await self._execute_on_target(edge.target_id, edge.source_id, message, shared_state, ctx) return True return False @@ -159,8 +179,20 @@ class FanInEdgeRunner(EdgeRunner): self._buffer.clear() # Send aggregated data to target aggregated_data = [msg.data for msg in messages_to_send] + + # Collect all trace contexts and source span IDs for fan-in linking + trace_contexts = [msg.trace_context for msg in messages_to_send if msg.trace_context] + source_span_ids = [msg.source_span_id for msg in messages_to_send if msg.source_span_id] + + # Create a new Message object for the aggregated data + aggregated_message = Message( + data=aggregated_data, + source_id=self._edge_group.__class__.__name__, + trace_contexts=trace_contexts, + source_span_ids=source_span_ids, + ) await self._execute_on_target( - self._edges[0].target_id, self._edge_group.__class__.__name__, aggregated_data, shared_state, ctx + self._edges[0].target_id, self._edge_group.__class__.__name__, aggregated_message, shared_state, ctx ) return True diff --git a/python/packages/workflow/agent_framework_workflow/_executor.py b/python/packages/workflow/agent_framework_workflow/_executor.py index daaac755bf..b466274ab1 100644 --- a/python/packages/workflow/agent_framework_workflow/_executor.py +++ b/python/packages/workflow/agent_framework_workflow/_executor.py @@ -78,27 +78,44 @@ class Executor(AFBaseModel): Returns: An awaitable that resolves to the result of the execution. """ + # Create processing span for tracing (gracefully handles disabled tracing) + from ._telemetry import workflow_tracer + + source_trace_contexts = getattr(context, "_trace_contexts", None) + source_span_ids = getattr(context, "_source_span_ids", None) + # Handle case where Message wrapper is passed instead of raw data + from ._runner_context import Message - # Lazy registration for SubWorkflowRequestInfo if we have interceptors - if self._request_interceptors and message.__class__.__name__ == "SubWorkflowRequestInfo": - # Directly handle SubWorkflowRequestInfo + if isinstance(message, Message): + message = message.data + + with workflow_tracer.create_processing_span( + self.id, + self.__class__.__name__, + type(message).__name__, + source_trace_contexts=source_trace_contexts, + source_span_ids=source_span_ids, + ): + # Lazy registration for SubWorkflowRequestInfo if we have interceptors + if self._request_interceptors and message.__class__.__name__ == "SubWorkflowRequestInfo": + # Directly handle SubWorkflowRequestInfo + await context.add_event(ExecutorInvokeEvent(self.id)) + await self._handle_sub_workflow_request(message, context) + await context.add_event(ExecutorCompletedEvent(self.id)) + return + + handler: Callable[[Any, WorkflowContext[Any]], Any] | None = None + for message_type in self._handlers: + if is_instance_of(message, message_type): + handler = self._handlers[message_type] + break + + if handler is None: + raise RuntimeError(f"Executor {self.__class__.__name__} cannot handle message of type {type(message)}.") await context.add_event(ExecutorInvokeEvent(self.id)) - await self._handle_sub_workflow_request(message, context) + await handler(message, context) await context.add_event(ExecutorCompletedEvent(self.id)) - return - - handler: Callable[[Any, WorkflowContext[Any]], Any] | None = None - for message_type in self._handlers: - if is_instance_of(message, message_type): - handler = self._handlers[message_type] - break - - if handler is None: - raise RuntimeError(f"Executor {self.__class__.__name__} cannot handle message of type {type(message)}.") - await context.add_event(ExecutorInvokeEvent(self.id)) - await handler(message, context) - await context.add_event(ExecutorCompletedEvent(self.id)) def _discover_handlers(self) -> None: """Discover message handlers and request interceptors in the executor class.""" @@ -196,17 +213,13 @@ class Executor(AFBaseModel): if correlated_response.is_handled: # Send response back to sub-workflow - from ._runner_context import Message - - response_message = Message( - source_id=self.id, - target_id=request.sub_workflow_id, - data=SubWorkflowResponse( + await ctx.send_message( + SubWorkflowResponse( request_id=request.request_id, data=correlated_response.data, ), + target_id=request.sub_workflow_id, ) - await ctx.send_message(response_message) else: # Forward WITH CONTEXT PRESERVED # Update the data if interceptor provided a modified request @@ -214,13 +227,7 @@ class Executor(AFBaseModel): request.data = correlated_response.forward_request # Send the inner request to RequestInfoExecutor to create external request - from ._runner_context import Message - - forward_message = Message( - source_id=self.id, - data=request, - ) - await ctx.send_message(forward_message) + await ctx.send_message(request) else: # Legacy support: direct return means handled await ctx.send_message( @@ -234,10 +241,7 @@ class Executor(AFBaseModel): # No interceptor found - forward inner request to RequestInfoExecutor # This sends the original request to RequestInfoExecutor - from ._runner_context import Message - - passthrough_message = Message(source_id=self.id, data=request.data) - await ctx.send_message(passthrough_message) + await ctx.send_message(request.data) def can_handle(self, message: Any) -> bool: """Check if the executor can handle a given message type. @@ -357,7 +361,7 @@ def handler( return await func(self, message, ctx) # Preserve the original function signature for introspection during validation - with contextlib.suppress(Exception): + with contextlib.suppress(AttributeError, TypeError): wrapper.__signature__ = sig # type: ignore[attr-defined] wrapper._handler_spec = { # type: ignore @@ -806,15 +810,19 @@ class WorkflowExecutor(Executor): are intercepted by parent workflows. """ - def __init__(self, workflow: "Workflow", id: str | None = None): + workflow: "Workflow" = Field(description="The workflow to execute as a sub-workflow") + + def __init__(self, workflow: "Workflow", id: str | None = None, **kwargs: Any): """Initialize the WorkflowExecutor. Args: workflow: The workflow to execute as a sub-workflow. id: Optional unique identifier for this executor. + **kwargs: Additional keyword arguments passed to the parent constructor. """ - super().__init__(id) - self._workflow = workflow + kwargs.update({"workflow": workflow}) + super().__init__(id, **kwargs) + # Track pending external responses by request_id self._pending_responses: dict[str, Any] = {} # request_id -> response_data # Track workflow state for proper resumption - support multiple concurrent requests @@ -846,7 +854,7 @@ class WorkflowExecutor(Executor): try: # Run the sub-workflow and collect all events - events = [event async for event in self._workflow.run_streaming(input_data)] + events = [event async for event in self.workflow.run_streaming(input_data)] # Count requests and initialize response tracking request_count = 0 @@ -932,7 +940,7 @@ class WorkflowExecutor(Executor): responses_to_send = dict(self._collected_responses) self._collected_responses.clear() # Clear for next batch - result_events = [event async for event in self._workflow.send_responses_streaming(responses_to_send)] + result_events = [event async for event in self.workflow.send_responses_streaming(responses_to_send)] # Process the result events new_request_count = 0 diff --git a/python/packages/workflow/agent_framework_workflow/_runner.py b/python/packages/workflow/agent_framework_workflow/_runner.py index 037430ed25..09921299b0 100644 --- a/python/packages/workflow/agent_framework_workflow/_runner.py +++ b/python/packages/workflow/agent_framework_workflow/_runner.py @@ -179,7 +179,16 @@ class Runner: f"from sub-workflow '{sub_request.sub_workflow_id}' " f"to executor '{executor.id}' for interception." ) - await executor.execute(sub_request, self._ctx) # type: ignore[arg-type] + # Create WorkflowContext with trace context from message + workflow_ctx: WorkflowContext[Any] = WorkflowContext( + executor.id, + [message.source_id], + self._shared_state, + self._ctx, + trace_contexts=[message.trace_context] if message.trace_context else None, + source_span_ids=[message.source_span_id] if message.source_span_id else None, + ) + await executor.execute(sub_request, workflow_ctx) interceptor_found = True break if interceptor_found: @@ -192,18 +201,20 @@ class Runner: request_info_executor = self._find_request_info_executor() if request_info_executor: - workflow_ctx: WorkflowContext[None] = WorkflowContext( + request_info_workflow_ctx: WorkflowContext[None] = WorkflowContext( request_info_executor.id, - ["Runner"], + [message.source_id], self._shared_state, self._ctx, + trace_contexts=[message.trace_context] if message.trace_context else None, + source_span_ids=[message.source_span_id] if message.source_span_id else None, ) logger.info( f"Sending sub-workflow request of type '{sub_request.data.__class__.__name__}' " f"from sub-workflow '{sub_request.sub_workflow_id}' to RequestInfoExecutor " f"'{request_info_executor.id}'" ) - await request_info_executor.execute(sub_request, workflow_ctx) + await request_info_executor.execute(sub_request, request_info_workflow_ctx) else: logger.warning( f"Sub-workflow request of type '{sub_request.data.__class__.__name__}' " diff --git a/python/packages/workflow/agent_framework_workflow/_runner_context.py b/python/packages/workflow/agent_framework_workflow/_runner_context.py index e6f1fab27a..990c70680d 100644 --- a/python/packages/workflow/agent_framework_workflow/_runner_context.py +++ b/python/packages/workflow/agent_framework_workflow/_runner_context.py @@ -24,6 +24,22 @@ class Message: source_id: str target_id: str | None = None + # OpenTelemetry trace context fields for message propagation + # These are plural to support fan-in scenarios where multiple messages are aggregated + trace_contexts: list[dict[str, str]] | None = None # W3C Trace Context headers from multiple sources + source_span_ids: list[str] | None = None # Publishing span IDs for linking from multiple sources + + # Backward compatibility properties + @property + def trace_context(self) -> dict[str, str] | None: + """Get the first trace context for backward compatibility.""" + return self.trace_contexts[0] if self.trace_contexts else None + + @property + def source_span_id(self) -> str | None: + """Get the first source span ID for backward compatibility.""" + return self.source_span_ids[0] if self.source_span_ids else None + class CheckpointState(TypedDict): messages: dict[str, list[dict[str, Any]]] @@ -268,7 +284,14 @@ class InProcRunnerContext: serializable_messages: dict[str, list[dict[str, Any]]] = {} for source_id, message_list in self._messages.items(): serializable_messages[source_id] = [ - {"data": msg.data, "source_id": msg.source_id, "target_id": msg.target_id} for msg in message_list + { + "data": msg.data, + "source_id": msg.source_id, + "target_id": msg.target_id, + "trace_contexts": msg.trace_contexts, + "source_span_ids": msg.source_span_ids, + } + for msg in message_list ] return { "messages": serializable_messages, @@ -287,6 +310,8 @@ class InProcRunnerContext: data=msg.get("data"), source_id=msg.get("source_id", ""), target_id=msg.get("target_id"), + trace_contexts=msg.get("trace_contexts"), + source_span_ids=msg.get("source_span_ids"), ) for msg in message_list ] diff --git a/python/packages/workflow/agent_framework_workflow/_telemetry.py b/python/packages/workflow/agent_framework_workflow/_telemetry.py new file mode 100644 index 0000000000..fad5233a7f --- /dev/null +++ b/python/packages/workflow/agent_framework_workflow/_telemetry.py @@ -0,0 +1,213 @@ +# Copyright (c) Microsoft. All rights reserved. + +from typing import TYPE_CHECKING, Any, ClassVar + +from agent_framework._pydantic import AFBaseSettings +from opentelemetry.trace import Link, NoOpTracer, SpanKind, StatusCode, get_current_span, get_tracer +from opentelemetry.trace.span import SpanContext +from opentelemetry.util.types import Attributes + +if TYPE_CHECKING: + from ._workflow import Workflow + + +# Span name constants +_WORKFLOW_BUILD_SPAN = "workflow.build" +_WORKFLOW_RUN_SPAN = "workflow.run" +_EXECUTOR_PROCESS_SPAN = "executor.process" +_MESSAGE_SEND_SPAN = "message.send" + + +class WorkflowDiagnosticSettings(AFBaseSettings): + """Settings for workflow tracing diagnostics.""" + + env_prefix: ClassVar[str] = "AGENT_FRAMEWORK_WORKFLOW_" + enable_otel_diagnostics: bool = False + + @property + def ENABLED(self) -> bool: + return self.enable_otel_diagnostics + + +class WorkflowTracer: + """Central tracing coordinator for workflow system. + + Manages OpenTelemetry span creation and relationships for: + - Workflow build spans (workflow.build) + - Workflow execution spans (workflow.run) + - Executor processing spans (executor.process) + - Message sending spans (message.send) + + Implements span linking for causality without unwanted nesting. + """ + + def __init__(self) -> None: + self.settings = WorkflowDiagnosticSettings() + self.tracer = get_tracer("agent_framework") if self.settings.ENABLED else NoOpTracer() + + @property + def enabled(self) -> bool: + return self.settings.ENABLED + + def create_workflow_run_span(self, workflow: "Workflow") -> Any: + """Create a workflow execution span.""" + attributes: dict[str, str | int] = { + "workflow.id": workflow.id, + } + + return self.tracer.start_as_current_span(_WORKFLOW_RUN_SPAN, kind=SpanKind.INTERNAL, attributes=attributes) + + def create_workflow_build_span(self) -> Any: + """Create a workflow build span.""" + return self.tracer.start_as_current_span(_WORKFLOW_BUILD_SPAN, kind=SpanKind.INTERNAL) + + def create_processing_span( + self, + executor_id: str, + executor_type: str, + message_type: str, + source_trace_contexts: list[dict[str, str]] | None = None, + source_span_ids: list[str] | None = None, + ) -> Any: + """Create an executor processing span with optional links to source spans. + + Processing spans are created as children of the current workflow span and + linked (not nested) to the source publishing spans for causality tracking. + This supports multiple links for fan-in scenarios. + """ + # Create links to source spans for causality without nesting + links = [] + if source_trace_contexts and source_span_ids: + # Create links for all source spans (supporting fan-in with multiple sources) + for trace_context, span_id in zip(source_trace_contexts, source_span_ids, strict=False): + try: + # Extract trace and span IDs from the trace context + # This is a simplified approach - in production you'd want more robust parsing + traceparent = trace_context.get("traceparent", "") + if traceparent: + # traceparent format: "00-{trace_id}-{parent_span_id}-{trace_flags}" + parts = traceparent.split("-") + if len(parts) >= 3: + trace_id_hex = parts[1] + # Use the source_span_id that was saved from the publishing span + + # Create span context for linking + span_context = SpanContext( + trace_id=int(trace_id_hex, 16), + span_id=int(span_id, 16), + is_remote=True, + ) + links.append(Link(span_context)) + except (ValueError, TypeError, AttributeError): + # If linking fails, continue without link (graceful degradation) + pass + + return self.tracer.start_as_current_span( + _EXECUTOR_PROCESS_SPAN, + kind=SpanKind.INTERNAL, + attributes={ + "executor.id": executor_id, + "executor.type": executor_type, + "message.type": message_type, + }, + links=links, + ) + + def create_sending_span(self, message_type: str, target_executor_id: str | None = None) -> Any: + """Create a message send span. + + Sending spans are created as children of the current processing span + to track message emission for distributed tracing. + """ + attributes: dict[str, str] = { + "message.type": message_type, + } + if target_executor_id is not None: + attributes["message.destination_executor_id"] = target_executor_id + + return self.tracer.start_as_current_span( + _MESSAGE_SEND_SPAN, + kind=SpanKind.PRODUCER, + attributes=attributes, + ) + + def add_workflow_event(self, event_name: str, attributes: Attributes | None = None) -> None: + """Add an event to the current workflow span. + + Args: + event_name: Name of the event (e.g., "workflow.started", "workflow.completed") + attributes: Optional attributes to attach to the event + """ + span = get_current_span() + if span and span.is_recording(): + span.add_event(event_name, attributes) + + def add_workflow_error_event(self, error: Exception, attributes: Attributes | None = None) -> None: + """Add an error event to the current workflow span. + + Args: + error: The exception that occurred + attributes: Optional additional attributes to attach to the event + """ + span = get_current_span() + if span and span.is_recording(): + event_attributes: dict[str, str | bool | int | float] = { + "error.message": str(error), + "error.type": type(error).__name__, + } + if attributes: + # Safely merge attributes, ensuring type compatibility + for key, value in attributes.items(): + if isinstance(value, (str, bool, int, float)): + event_attributes[key] = value + span.add_event("workflow.error", event_attributes) + span.set_status(StatusCode.ERROR, str(error)) + + def set_workflow_build_span_attributes(self, workflow: "Workflow") -> None: + """Set workflow attributes on the current span. + + Args: + workflow: The workflow instance to extract attributes from + """ + span = get_current_span() + if span and span.is_recording(): + span.set_attributes({ + "workflow.id": workflow.id, + "workflow.definition": workflow.model_dump_json(), + }) + + def add_build_event(self, event_name: str, attributes: Attributes | None = None) -> None: + """Add an event to the current workflow build span. + + Args: + event_name: Name of the build event (e.g., "build.started", "build.validation_completed") + attributes: Optional attributes to attach to the event + """ + span = get_current_span() + if span and span.is_recording(): + span.add_event(event_name, attributes) + + def add_build_error_event(self, error: Exception, attributes: Attributes | None = None) -> None: + """Add an error event to the current workflow build span. + + Args: + error: The exception that occurred during build + attributes: Optional additional attributes to attach to the event + """ + span = get_current_span() + if span and span.is_recording(): + event_attributes: dict[str, str | bool | int | float] = { + "build.error.message": str(error), + "build.error.type": type(error).__name__, + } + if attributes: + # Safely merge attributes, ensuring type compatibility + for key, value in attributes.items(): + if isinstance(value, (str, bool, int, float)): + event_attributes[key] = value + span.add_event("build.error", event_attributes) + span.set_status(StatusCode.ERROR, str(error)) + + +# Global workflow tracer instance +workflow_tracer = WorkflowTracer() diff --git a/python/packages/workflow/agent_framework_workflow/_workflow.py b/python/packages/workflow/agent_framework_workflow/_workflow.py index f50ccc9e05..683697218d 100644 --- a/python/packages/workflow/agent_framework_workflow/_workflow.py +++ b/python/packages/workflow/agent_framework_workflow/_workflow.py @@ -4,7 +4,7 @@ import asyncio import logging import sys import uuid -from collections.abc import AsyncIterable, Callable, Sequence +from collections.abc import AsyncIterable, Awaitable, Callable, Sequence from typing import Any from agent_framework._pydantic import AFBaseModel @@ -88,7 +88,7 @@ class Workflow(AFBaseModel): max_iterations: int = Field( default=DEFAULT_MAX_ITERATIONS, description="Maximum number of iterations the workflow will run" ) - workflow_id: str = Field( + id: str = Field( default_factory=lambda: str(uuid.uuid4()), description="Unique identifier for this workflow instance" ) @@ -114,14 +114,14 @@ class Workflow(AFBaseModel): # Convert start_executor to string ID if it's an Executor instance start_executor_id = start_executor.id if isinstance(start_executor, Executor) else start_executor - workflow_id = str(uuid.uuid4()) + id = str(uuid.uuid4()) kwargs.update({ "edge_groups": edge_groups, "executors": executors, "start_executor_id": start_executor_id, "max_iterations": max_iterations, - "workflow_id": workflow_id, + "id": id, }) super().__init__(**kwargs) @@ -135,9 +135,39 @@ class Workflow(AFBaseModel): self._shared_state, runner_context, max_iterations=max_iterations, - workflow_id=workflow_id, + workflow_id=id, ) + def model_dump(self, **kwargs: Any) -> dict[str, Any]: + """Custom serialization that properly handles WorkflowExecutor nested workflows.""" + data = super().model_dump(**kwargs) + + # Ensure WorkflowExecutor instances have their workflow field serialized + if "executors" in data: + executors_data = data["executors"] + for executor_id, executor_data in executors_data.items(): + # Check if this is a WorkflowExecutor that might be missing its workflow field + if ( + isinstance(executor_data, dict) + and executor_data.get("type") == "WorkflowExecutor" + and "workflow" not in executor_data + ): + # Get the original executor object and serialize its workflow + original_executor = self.executors.get(executor_id) + if original_executor and hasattr(original_executor, "workflow"): + from ._executor import WorkflowExecutor + + if isinstance(original_executor, WorkflowExecutor): + executor_data["workflow"] = original_executor.workflow.model_dump(**kwargs) + + return data + + def model_dump_json(self, **kwargs: Any) -> str: + """Custom JSON serialization that properly handles WorkflowExecutor nested workflows.""" + import json + + return json.dumps(self.model_dump(**kwargs)) + def get_start_executor(self) -> Executor: """Get the starting executor of the workflow. @@ -150,6 +180,48 @@ class Workflow(AFBaseModel): """Get the list of executors in the workflow.""" return list(self.executors.values()) + async def _run_workflow_with_tracing( + self, initial_executor_fn: Callable[[], Awaitable[None]] | None = None, reset_context: bool = True + ) -> AsyncIterable[WorkflowEvent]: + """Private method to run workflow with proper tracing. + + All workflow entry points create a NEW workflow span. It is the responsibility + of external callers to maintain context across different workflow runs. + + Args: + initial_executor_fn: Optional function to execute initial executor + reset_context: Whether to reset the context for a new run + + Yields: + WorkflowEvent: The events generated during the workflow execution. + """ + # Import here to avoid circular imports + from ._telemetry import workflow_tracer + + # Create workflow span that encompasses the entire execution + with workflow_tracer.create_workflow_run_span(self): + try: + # Add workflow started event + workflow_tracer.add_workflow_event("workflow.started") + + # Reset context for a new run if supported + if reset_context: + self._runner.context.reset_for_new_run(self._shared_state) + + # Execute initial setup if provided + if initial_executor_fn: + await initial_executor_fn() + + # All executor executions happen within workflow span + async for event in self._runner.run_until_convergence(): + yield event + + # Success + workflow_tracer.add_workflow_event("workflow.completed") + except Exception as e: + workflow_tracer.add_workflow_error_event(e) + raise + async def run_streaming(self, message: Any) -> AsyncIterable[WorkflowEvent]: """Run the workflow with a starting message and stream events. @@ -159,22 +231,22 @@ class Workflow(AFBaseModel): Yields: WorkflowEvent: The events generated during the workflow execution. """ - # Reset context for a new run if supported - self._runner.context.reset_for_new_run(self._shared_state) - executor = self.get_start_executor() + async def initial_execution() -> None: + executor = self.get_start_executor() + await executor.execute( + message, + WorkflowContext( + executor.id, + [self.__class__.__name__], + self._shared_state, + self._runner.context, + trace_contexts=None, # No parent trace context for workflow start + source_span_ids=None, # No source span for workflow start + ), + ) - await executor.execute( - message, - WorkflowContext( - executor.id, - [self.__class__.__name__], - self._shared_state, - self._runner.context, - ), - ) - - async for event in self._runner.run_until_convergence(): + async for event in self._run_workflow_with_tracing(initial_executor_fn=initial_execution, reset_context=True): yield event async def run_streaming_from_checkpoint( @@ -199,41 +271,48 @@ class Workflow(AFBaseModel): ValueError: If neither checkpoint_storage is provided nor checkpointing is enabled. RuntimeError: If checkpoint restoration fails. """ - has_checkpointing = self._runner.context.has_checkpointing() - if not has_checkpointing and not checkpoint_storage: - raise ValueError( - "Cannot restore from checkpoint: either provide checkpoint_storage parameter " - "or build workflow with WorkflowBuilder.with_checkpointing(checkpoint_storage)." - ) + async def checkpoint_restoration() -> None: + has_checkpointing = self._runner.context.has_checkpointing() - if has_checkpointing: - # restore via Runner so shared state and iteration are synchronized - restored = await self._runner.restore_from_checkpoint(checkpoint_id) - else: - if checkpoint_storage is None: - raise ValueError("checkpoint_storage cannot be None.") - restored = await self._restore_from_external_checkpoint(checkpoint_id, checkpoint_storage) + if not has_checkpointing and not checkpoint_storage: + raise ValueError( + "Cannot restore from checkpoint: either provide checkpoint_storage parameter " + "or build workflow with WorkflowBuilder.with_checkpointing(checkpoint_storage)." + ) - if not restored: - raise RuntimeError(f"Failed to restore from checkpoint: {checkpoint_id}") + if has_checkpointing: + # restore via Runner so shared state and iteration are synchronized + restored = await self._runner.restore_from_checkpoint(checkpoint_id) + else: + if checkpoint_storage is None: + raise ValueError("checkpoint_storage cannot be None.") + restored = await self._restore_from_external_checkpoint(checkpoint_id, checkpoint_storage) - if responses: - request_info_executor = self._find_request_info_executor() - if request_info_executor: - for request_id, response_data in responses.items(): - await request_info_executor.handle_response( - response_data, - request_id, - WorkflowContext( - request_info_executor.id, - [self.__class__.__name__], - self._shared_state, - self._runner.context, - ), - ) + if not restored: + raise RuntimeError(f"Failed to restore from checkpoint: {checkpoint_id}") - async for event in self._runner.run_until_convergence(): + if responses: + request_info_executor = self._find_request_info_executor() + if request_info_executor: + for request_id, response_data in responses.items(): + await request_info_executor.handle_response( + response_data, + request_id, + WorkflowContext( + request_info_executor.id, + [self.__class__.__name__], + self._shared_state, + self._runner.context, + trace_contexts=None, # No parent trace context for new workflow span + source_span_ids=None, # No source span for response handling + ), + ) + + async for event in self._run_workflow_with_tracing( + initial_executor_fn=checkpoint_restoration, + reset_context=False, # Don't reset context when resuming from checkpoint + ): yield event async def send_responses_streaming(self, responses: dict[str, Any]) -> AsyncIterable[WorkflowEvent]: @@ -246,26 +325,35 @@ class Workflow(AFBaseModel): Yields: WorkflowEvent: The events generated during the workflow execution after sending the responses. """ - request_info_executor = self._find_request_info_executor() - if not request_info_executor: - raise ValueError("No RequestInfoExecutor found in workflow.") - async def _handle_response(response: Any, request_id: str) -> None: - """Handle the response from the RequestInfoExecutor.""" - await request_info_executor.handle_response( - response, - request_id, - WorkflowContext( - request_info_executor.id, - [self.__class__.__name__], - self._shared_state, - self._runner.context, - ), - ) + async def send_responses() -> None: + request_info_executor = self._find_request_info_executor() + if not request_info_executor: + raise ValueError("No RequestInfoExecutor found in workflow.") - await asyncio.gather(*[_handle_response(response, request_id) for request_id, response in responses.items()]) + async def _handle_response(response: Any, request_id: str) -> None: + """Handle the response from the RequestInfoExecutor.""" + await request_info_executor.handle_response( + response, + request_id, + WorkflowContext( + request_info_executor.id, + [self.__class__.__name__], + self._shared_state, + self._runner.context, + trace_contexts=None, # No parent trace context for new workflow span + source_span_ids=None, # No source span for response handling + ), + ) - async for event in self._runner.run_until_convergence(): + await asyncio.gather(*[ + _handle_response(response, request_id) for request_id, response in responses.items() + ]) + + async for event in self._run_workflow_with_tracing( + initial_executor_fn=send_responses, + reset_context=False, # Don't reset context when sending responses + ): yield event async def run(self, message: Any) -> WorkflowRunResult: @@ -437,7 +525,13 @@ class Workflow(AFBaseModel): from ._runner_context import Message as _Msg await self._runner.context.send_message( - _Msg(data=msg_data.get("data"), source_id=source_id, target_id=target_id) + _Msg( + data=msg_data.get("data"), + source_id=source_id, + target_id=target_id, + trace_contexts=msg_data.get("trace_contexts"), + source_span_ids=msg_data.get("source_span_ids"), + ) ) @@ -652,14 +746,42 @@ class WorkflowBuilder: WorkflowValidationError: If workflow validation fails (includes EdgeDuplicationError, TypeCompatibilityError, and GraphConnectivityError subclasses). """ - if not self._start_executor: - raise ValueError("Starting executor must be set using set_start_executor before building the workflow.") + # Import here to avoid circular imports + from ._telemetry import workflow_tracer - validate_workflow_graph(self._edge_groups, self._executors, self._start_executor) + # Create workflow build span that includes validation and workflow creation + with workflow_tracer.create_workflow_build_span(): + try: + # Add workflow build started event + workflow_tracer.add_build_event("build.started") - context = InProcRunnerContext(self._checkpoint_storage) + if not self._start_executor: + raise ValueError( + "Starting executor must be set using set_start_executor before building the workflow." + ) - return Workflow(self._edge_groups, self._executors, self._start_executor, context, self._max_iterations) + # Perform validation before creating the workflow + validate_workflow_graph(self._edge_groups, self._executors, self._start_executor) + # Add validation completed event + workflow_tracer.add_build_event("build.validation_completed") -# endregion + context = InProcRunnerContext(self._checkpoint_storage) + + # Create workflow instance after validation + workflow = Workflow( + self._edge_groups, self._executors, self._start_executor, context, self._max_iterations + ) + + # Set workflow attributes on the span + workflow_tracer.set_workflow_build_span_attributes(workflow) + + # Add workflow build completed event + workflow_tracer.add_build_event("build.completed") + + return workflow + + except Exception as e: + # The method already includes sufficient error info (error.message, error.type) + workflow_tracer.add_build_error_event(e) + raise diff --git a/python/packages/workflow/agent_framework_workflow/_workflow_context.py b/python/packages/workflow/agent_framework_workflow/_workflow_context.py index 255fdfa790..219f164966 100644 --- a/python/packages/workflow/agent_framework_workflow/_workflow_context.py +++ b/python/packages/workflow/agent_framework_workflow/_workflow_context.py @@ -2,9 +2,12 @@ from typing import Any, Generic, TypeVar +from opentelemetry.propagate import inject + from ._events import WorkflowEvent from ._runner_context import Message, RunnerContext from ._shared_state import SharedState +from ._telemetry import workflow_tracer T_Out = TypeVar("T_Out") @@ -22,6 +25,8 @@ class WorkflowContext(Generic[T_Out]): source_executor_ids: list[str], shared_state: SharedState, runner_context: RunnerContext, + trace_contexts: list[dict[str, str]] | None = None, + source_span_ids: list[str] | None = None, ): """Initialize the executor context with the given workflow context. @@ -32,12 +37,18 @@ class WorkflowContext(Generic[T_Out]): messages to the same executor. shared_state: The shared state for the workflow. runner_context: The runner context that provides methods to send messages and events. + trace_contexts: Optional trace contexts from multiple sources for OpenTelemetry propagation. + source_span_ids: Optional source span IDs from multiple sources for linking (not for nesting). """ self._executor_id = executor_id self._source_executor_ids = source_executor_ids self._runner_context = runner_context self._shared_state = shared_state + # Store trace contexts and source span IDs for linking (supporting multiple sources) + self._trace_contexts = trace_contexts or [] + self._source_span_ids = source_span_ids or [] + if not self._source_executor_ids: raise ValueError("source_executor_ids cannot be empty. At least one source executor ID is required.") @@ -49,13 +60,20 @@ class WorkflowContext(Generic[T_Out]): target_id: The ID of the target executor to send the message to. If None, the message will be sent to all target executors. """ - await self._runner_context.send_message( - Message( - data=message, - source_id=self._executor_id, - target_id=target_id, - ) - ) + # Create publishing span (inherits current trace context automatically) + with workflow_tracer.create_sending_span(type(message).__name__, target_id) as span: + # Create Message wrapper + msg = Message(data=message, source_id=self._executor_id, target_id=target_id) + + # Inject current trace context if tracing enabled + if workflow_tracer.enabled and span and span.is_recording(): + trace_context: dict[str, str] = {} + inject(trace_context) # Inject current trace context for message propagation + + msg.trace_contexts = [trace_context] + msg.source_span_ids = [format(span.get_span_context().span_id, "016x")] + + await self._runner_context.send_message(msg) async def add_event(self, event: WorkflowEvent) -> None: """Add an event to the workflow context.""" diff --git a/python/packages/workflow/tests/test_serialization.py b/python/packages/workflow/tests/test_serialization.py index 02c405662d..46de4d3f04 100644 --- a/python/packages/workflow/tests/test_serialization.py +++ b/python/packages/workflow/tests/test_serialization.py @@ -15,6 +15,9 @@ from agent_framework_workflow._edge import ( SwitchCaseEdgeGroupCase, SwitchCaseEdgeGroupDefault, ) +from agent_framework_workflow._executor import ( + WorkflowExecutor, +) class SampleExecutor(Executor): @@ -398,6 +401,110 @@ class TestSerializationWorkflowClasses: "JSON SwitchCaseEdgeGroup edges should not have condition_name" ) + def test_nested_workflow_executor_serialization(self) -> None: + """Test complete serialization of deeply nested WorkflowExecutors (subworkflows within subworkflows). + + This test verifies that nested WorkflowExecutor objects are fully serialized with their + complete workflow structures, including deeply nested workflows and all their executors. + """ + # Create innermost workflow + inner_executor = SampleExecutor(id="inner-exec") + inner_workflow = WorkflowBuilder().set_start_executor(inner_executor).set_max_iterations(10).build() + + # Create middle workflow with WorkflowExecutor + inner_workflow_executor = WorkflowExecutor(workflow=inner_workflow, id="inner-workflow-exec") + middle_executor = SampleExecutor(id="middle-exec") + middle_workflow = ( + WorkflowBuilder() + .set_start_executor(middle_executor) + .add_edge(middle_executor, inner_workflow_executor) + .set_max_iterations(20) + .build() + ) + + # Create outer workflow with nested WorkflowExecutor + middle_workflow_executor = WorkflowExecutor(workflow=middle_workflow, id="middle-workflow-exec") + outer_executor = SampleExecutor(id="outer-exec") + outer_workflow = ( + WorkflowBuilder() + .set_start_executor(outer_executor) + .add_edge(outer_executor, middle_workflow_executor) + .set_max_iterations(30) + .build() + ) + + # Test serialization of the nested structure + data = outer_workflow.model_dump() + + # Verify outer structure + assert data["start_executor_id"] == "outer-exec" + assert data["max_iterations"] == 30 + assert "outer-exec" in data["executors"] + assert "middle-workflow-exec" in data["executors"] + + # Verify middle WorkflowExecutor is present with full nested workflow serialization + middle_exec_data = data["executors"]["middle-workflow-exec"] + assert middle_exec_data["type"] == "WorkflowExecutor" + assert middle_exec_data["id"] == "middle-workflow-exec" + + # Verify the nested workflow is fully serialized + assert "workflow" in middle_exec_data, "WorkflowExecutor should include nested workflow in serialization" + middle_workflow_data = middle_exec_data["workflow"] + assert "start_executor_id" in middle_workflow_data + assert "executors" in middle_workflow_data + assert "max_iterations" in middle_workflow_data + assert middle_workflow_data["start_executor_id"] == "middle-exec" + assert middle_workflow_data["max_iterations"] == 20 + + # Verify the deeply nested executors are present + assert "middle-exec" in middle_workflow_data["executors"] + assert "inner-workflow-exec" in middle_workflow_data["executors"] + + # Verify the innermost WorkflowExecutor is also fully serialized + inner_workflow_exec_data = middle_workflow_data["executors"]["inner-workflow-exec"] + assert inner_workflow_exec_data["type"] == "WorkflowExecutor" + assert "workflow" in inner_workflow_exec_data, "Deeply nested WorkflowExecutor should also include its workflow" + innermost_workflow_data = inner_workflow_exec_data["workflow"] + assert "start_executor_id" in innermost_workflow_data + assert "executors" in innermost_workflow_data + assert "max_iterations" in innermost_workflow_data + assert innermost_workflow_data["start_executor_id"] == "inner-exec" + assert innermost_workflow_data["max_iterations"] == 10 + assert "inner-exec" in innermost_workflow_data["executors"] + + # Test JSON serialization preserves the complete nested structure + json_str = outer_workflow.model_dump_json() + parsed = json.loads(json_str) + + # Verify the complete structure is preserved in JSON + middle_exec_json = parsed["executors"]["middle-workflow-exec"] + assert middle_exec_json["type"] == "WorkflowExecutor" + assert middle_exec_json["id"] == "middle-workflow-exec" + + # Verify nested workflow is present in JSON + assert "workflow" in middle_exec_json, "JSON serialization should include nested workflow" + middle_workflow_json = middle_exec_json["workflow"] + assert middle_workflow_json["start_executor_id"] == "middle-exec" + assert middle_workflow_json["max_iterations"] == 20 + assert "middle-exec" in middle_workflow_json["executors"] + assert "inner-workflow-exec" in middle_workflow_json["executors"] + + # Verify deeply nested structure in JSON + inner_workflow_exec_json = middle_workflow_json["executors"]["inner-workflow-exec"] + assert inner_workflow_exec_json["type"] == "WorkflowExecutor" + assert "workflow" in inner_workflow_exec_json, "Deeply nested WorkflowExecutor should be in JSON" + innermost_workflow_json = inner_workflow_exec_json["workflow"] + assert innermost_workflow_json["start_executor_id"] == "inner-exec" + assert innermost_workflow_json["max_iterations"] == 10 + assert "inner-exec" in innermost_workflow_json["executors"] + + # Test that WorkflowExecutor also serializes correctly when accessed directly + direct_middle_data = middle_workflow_executor.model_dump() + assert "workflow" in direct_middle_data + assert direct_middle_data["type"] == "WorkflowExecutor" + assert "executors" in direct_middle_data["workflow"] + assert "inner-workflow-exec" in direct_middle_data["workflow"]["executors"] + def test_switch_case_edge_group_serialization_with_named_condition(self) -> None: """Test that SwitchCaseEdgeGroup with named condition function serializes condition_name correctly.""" @@ -440,7 +547,7 @@ class TestSerializationWorkflowClasses: assert "executors" in data assert "start_executor_id" in data assert "max_iterations" in data - assert "workflow_id" in data + assert "id" in data assert data["start_executor_id"] == "executor1" assert "executor1" in data["executors"] diff --git a/python/packages/workflow/tests/test_tracing.py b/python/packages/workflow/tests/test_tracing.py new file mode 100644 index 0000000000..5dda235393 --- /dev/null +++ b/python/packages/workflow/tests/test_tracing.py @@ -0,0 +1,532 @@ +# Copyright (c) Microsoft. All rights reserved. + +import os +from collections.abc import Generator +from typing import Any, cast + +import pytest +from opentelemetry import trace +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from agent_framework_workflow import WorkflowBuilder +from agent_framework_workflow._executor import Executor, handler +from agent_framework_workflow._runner_context import InProcRunnerContext, Message +from agent_framework_workflow._shared_state import SharedState +from agent_framework_workflow._telemetry import WorkflowTracer, workflow_tracer +from agent_framework_workflow._workflow import Workflow +from agent_framework_workflow._workflow_context import WorkflowContext + + +@pytest.fixture +def tracing_enabled() -> Generator[None, None, None]: + """Enable tracing for tests.""" + original_value = os.environ.get("AGENT_FRAMEWORK_WORKFLOW_ENABLE_OTEL_DIAGNOSTICS") + os.environ["AGENT_FRAMEWORK_WORKFLOW_ENABLE_OTEL_DIAGNOSTICS"] = "true" + + # Force reload the settings to pick up the environment variable + from agent_framework_workflow._telemetry import WorkflowDiagnosticSettings + + workflow_tracer.settings = WorkflowDiagnosticSettings() + + yield + + # Restore original value + if original_value is None: + os.environ.pop("AGENT_FRAMEWORK_WORKFLOW_ENABLE_OTEL_DIAGNOSTICS", None) + else: + os.environ["AGENT_FRAMEWORK_WORKFLOW_ENABLE_OTEL_DIAGNOSTICS"] = original_value + + # Reload settings again + workflow_tracer.settings = WorkflowDiagnosticSettings() + + +@pytest.fixture +def span_exporter(tracing_enabled: Any) -> Generator[InMemorySpanExporter, None, None]: + """Set up OpenTelemetry test infrastructure.""" + # Use the built-in InMemorySpanExporter for better compatibility + exporter = InMemorySpanExporter() + tracer_provider = TracerProvider() + tracer_provider.add_span_processor(SimpleSpanProcessor(exporter)) + + # Store original tracer + original_tracer = workflow_tracer.tracer + + # Set up our test tracer + workflow_tracer.tracer = tracer_provider.get_tracer("agent_framework") + + yield exporter + + # Clean up + exporter.clear() + workflow_tracer.tracer = original_tracer + + +class MockExecutor(Executor): + """Mock executor for testing.""" + + def __init__(self, id: str = "mock_executor") -> None: + super().__init__(id=id) + # Use private field to avoid Pydantic validation + self._processed_messages: list[str] = [] + + @handler + async def handle_message(self, message: str, ctx: WorkflowContext[str]) -> None: + """Handle string messages.""" + self._processed_messages.append(message) + await ctx.send_message(f"processed: {message}") + + @property + def processed_messages(self) -> list[str]: + """Access to processed messages for testing.""" + return self._processed_messages + + +class SecondExecutor(Executor): + """Second executor for testing message chains.""" + + def __init__(self, id: str = "second_executor") -> None: + super().__init__(id=id) + # Use private field to avoid Pydantic validation + self._processed_messages: list[str] = [] + + @handler + async def handle_message(self, message: str, ctx: WorkflowContext[None]) -> None: + """Handle string messages.""" + self._processed_messages.append(message) + + @property + def processed_messages(self) -> list[str]: + """Access to processed messages for testing.""" + return self._processed_messages + + +class ProcessingExecutor(Executor): + """Executor that processes and forwards messages with a custom prefix.""" + + def __init__(self, id: str, prefix: str = "processed") -> None: + super().__init__(id=id) + # Use private field to avoid Pydantic validation + self._processed_messages: list[str] = [] + self._prefix = prefix + + @handler + async def handle_message(self, message: str, ctx: WorkflowContext[str]) -> None: + """Handle string messages and send them forward with prefix.""" + self._processed_messages.append(message) + await ctx.send_message(f"{self._prefix}: {message}") + + @property + def processed_messages(self) -> list[str]: + return self._processed_messages + + +class FanInAggregator(Executor): + """Fan-in aggregator that expects a list of inputs.""" + + def __init__(self, id: str = "aggregator") -> None: + super().__init__(id=id) + # Use private field to avoid Pydantic validation + self._processed_messages: list[Any] = [] + + @handler + async def handle_aggregated_data(self, messages: list[str], ctx: WorkflowContext[None]) -> None: + # Process aggregated messages from fan-in + aggregated = f"aggregated: {', '.join(messages)}" + self._processed_messages.append(aggregated) + + @property + def processed_messages(self) -> list[Any]: + """Access to processed messages for testing.""" + return self._processed_messages + + +@pytest.mark.asyncio +async def test_workflow_tracer_configuration() -> None: + """Test that workflow tracer can be enabled and disabled.""" + # Test disabled by default + tracer = WorkflowTracer() + assert not tracer.enabled + + # Test enabled with environment variable + original_value = os.environ.get("AGENT_FRAMEWORK_WORKFLOW_ENABLE_OTEL_DIAGNOSTICS") + os.environ["AGENT_FRAMEWORK_WORKFLOW_ENABLE_OTEL_DIAGNOSTICS"] = "true" + + # Force reload the settings to pick up the environment variable + from agent_framework_workflow._telemetry import WorkflowDiagnosticSettings + + tracer.settings = WorkflowDiagnosticSettings() + + assert tracer.enabled + + # Restore original value + if original_value is None: + os.environ.pop("AGENT_FRAMEWORK_WORKFLOW_ENABLE_OTEL_DIAGNOSTICS", None) + else: + os.environ["AGENT_FRAMEWORK_WORKFLOW_ENABLE_OTEL_DIAGNOSTICS"] = original_value + + # Reload settings again + tracer.settings = WorkflowDiagnosticSettings() + + +@pytest.mark.asyncio +async def test_span_creation_and_attributes(tracing_enabled: Any, span_exporter: InMemorySpanExporter) -> None: + """Test creation and attributes of all span types (workflow, processing, sending).""" + # Create a mock workflow object + mock_workflow = cast( + Workflow, + type( + "MockWorkflow", + (), + { + "id": "test-workflow-123", + "max_iterations": 100, + "model_dump_json": lambda self: '{"id": "test-workflow-123", "type": "mock"}', + }, + )(), + ) + + # Test all span types in nested context + with workflow_tracer.create_workflow_run_span(mock_workflow) as workflow_span: + workflow_tracer.add_workflow_event("workflow.started") + + with ( + workflow_tracer.create_processing_span("executor-456", "TestExecutor", "TestMessage") as processing_span, + workflow_tracer.create_sending_span("ResponseMessage", "target-789") as sending_span, + ): + # Verify all spans are recording + assert workflow_span is not None and workflow_span.is_recording() + assert processing_span is not None and processing_span.is_recording() + assert sending_span is not None and sending_span.is_recording() + + spans = span_exporter.get_finished_spans() + assert len(spans) == 3 + + # Check workflow span + workflow_span = next(s for s in spans if s.name == "workflow.run") + assert workflow_span.kind == trace.SpanKind.INTERNAL + assert workflow_span.attributes is not None + assert workflow_span.attributes.get("workflow.id") == "test-workflow-123" + assert workflow_span.events is not None + event_names = [event.name for event in workflow_span.events] + assert "workflow.started" in event_names + + # Check processing span + processing_span = next(s for s in spans if s.name == "executor.process") + assert processing_span.kind == trace.SpanKind.INTERNAL + assert processing_span.attributes is not None + assert processing_span.attributes.get("executor.id") == "executor-456" + assert processing_span.attributes.get("executor.type") == "TestExecutor" + assert processing_span.attributes.get("message.type") == "TestMessage" + + # Check sending span + sending_span = next(s for s in spans if s.name == "message.send") + assert sending_span.kind == trace.SpanKind.PRODUCER + assert sending_span.attributes is not None + assert sending_span.attributes.get("message.type") == "ResponseMessage" + assert sending_span.attributes.get("message.destination_executor_id") == "target-789" + + +@pytest.mark.asyncio +async def test_trace_context_handling(tracing_enabled: Any, span_exporter: InMemorySpanExporter) -> None: + """Test trace context propagation and handling in messages and executors.""" + shared_state = SharedState() + ctx = InProcRunnerContext() + executor = MockExecutor("test-executor") + + # Test trace context propagation in messages + workflow_ctx: WorkflowContext[str] = WorkflowContext( + "test-executor", + ["source"], + shared_state, + ctx, + trace_contexts=[{"traceparent": "00-12345678901234567890123456789012-1234567890123456-01"}], + source_span_ids=["1234567890123456"], + ) + + # Send a message (this should create a sending span and propagate trace context) + await workflow_ctx.send_message("test message") + + # Check that message was created with trace context + messages = await ctx.drain_messages() + assert len(messages) == 1 + message_list = list(messages.values())[0] + assert len(message_list) == 1 + message = message_list[0] + assert message.trace_context is not None + assert message.source_span_id is not None + + # Test executor trace context handling + await executor.execute("test message", workflow_ctx) + + # Check that spans were created with proper attributes + spans = span_exporter.get_finished_spans() + processing_spans = [s for s in spans if s.name == "executor.process"] + sending_spans = [s for s in spans if s.name == "message.send"] + + assert len(processing_spans) >= 1 + assert len(sending_spans) >= 1 + + # Verify processing span attributes + processing_span = processing_spans[0] + assert processing_span.attributes is not None + assert processing_span.attributes.get("executor.id") == "test-executor" + assert processing_span.attributes.get("executor.type") == "MockExecutor" + assert processing_span.attributes.get("message.type") == "str" + + +@pytest.mark.asyncio +async def test_trace_context_disabled_when_tracing_disabled() -> None: + """Test that no trace context is added when tracing is disabled.""" + # Tracing should be disabled by default + shared_state = SharedState() + ctx = InProcRunnerContext() + + workflow_ctx: WorkflowContext[str] = WorkflowContext( + "test-executor", + ["source"], + shared_state, + ctx, + ) + + # Send a message + await workflow_ctx.send_message("test message") + + # Check that message was created without trace context + messages = await ctx.drain_messages() + message = list(messages.values())[0][0] + + # When tracing is disabled, trace_context should be None + assert message.trace_context is None + assert message.source_span_id is None + + +@pytest.mark.asyncio +async def test_end_to_end_workflow_tracing(tracing_enabled: Any, span_exporter: InMemorySpanExporter) -> None: + """Test end-to-end tracing including workflow build, execution, and span linking with fan-in edges.""" + # Create executors for fan-in scenario + executor1 = MockExecutor("executor1") + executor2 = ProcessingExecutor("executor2", "second") + executor3 = ProcessingExecutor("executor3", "third") + aggregator = FanInAggregator("aggregator") + + # Create workflow with fan-in: executor1 -> [executor2, executor3] -> aggregator + workflow = ( + WorkflowBuilder() + .set_start_executor(executor1) + .add_fan_out_edges(executor1, [executor2, executor3]) + .add_fan_in_edges([executor2, executor3], aggregator) + .build() + ) + + # Verify build span was created + build_spans = [s for s in span_exporter.get_finished_spans() if s.name == "workflow.build"] + assert len(build_spans) == 1 + + build_span = build_spans[0] + assert build_span.attributes is not None + assert build_span.attributes.get("workflow.id") == workflow.id + assert build_span.attributes.get("workflow.definition") is not None + definition = build_span.attributes.get("workflow.definition") + assert definition == workflow.model_dump_json() + + # Check build events + assert build_span.events is not None + build_event_names = [event.name for event in build_span.events] + assert "build.started" in build_event_names + assert "build.validation_completed" in build_event_names + assert "build.completed" in build_event_names + + # Clear spans to separate build from run tracing + span_exporter.clear() + + # Run workflow (this should create run spans) + events = [] + async for event in workflow.run_streaming("test input"): + events.append(event) + + # Verify workflow executed correctly + assert len(executor1.processed_messages) == 1 + assert executor1.processed_messages[0] == "test input" + assert len(executor2.processed_messages) == 1 + assert executor2.processed_messages[0] == "processed: test input" + assert len(executor3.processed_messages) == 1 + assert executor3.processed_messages[0] == "processed: test input" # executor3 receives from executor1 via fan-out + assert len(aggregator.processed_messages) == 1 + # The aggregator should receive both processed messages from executor2 and executor3 + aggregated_msg = aggregator.processed_messages[0] + assert "second: processed: test input" in aggregated_msg + assert "third: processed: test input" in aggregated_msg + + # Check run spans (build spans should not be present after clear) + spans = span_exporter.get_finished_spans() + + # Should have workflow span, processing spans, and sending spans + workflow_spans = [s for s in spans if s.name == "workflow.run"] + processing_spans = [s for s in spans if s.name == "executor.process"] + sending_spans = [s for s in spans if s.name == "message.send"] + build_spans_after_run = [s for s in spans if s.name == "workflow.build"] + + assert len(workflow_spans) == 1 + assert len(processing_spans) >= 4 # executor1, executor2, executor3, aggregator + assert len(sending_spans) >= 3 # Messages sent between executors + assert len(build_spans_after_run) == 0 # No build spans should be present after clear + + # Verify workflow span events + workflow_span = workflow_spans[0] + assert workflow_span.events is not None + event_names = [event.name for event in workflow_span.events] + assert "workflow.started" in event_names + assert "workflow.completed" in event_names + + # Test fan-in span linking: find the aggregator's processing span + aggregator_spans = [s for s in processing_spans if s.attributes and s.attributes.get("executor.id") == "aggregator"] + assert len(aggregator_spans) == 1 + + aggregator_span = aggregator_spans[0] + # The aggregator span should have links to the source spans (from executor2 and executor3) + # This tests that FanInEdgeRunner properly handles multiple trace contexts and span IDs + assert aggregator_span.links is not None + + # Find the sending spans from executor2 and executor3 by checking parent relationships + executor2_processing_spans = [ + s for s in processing_spans if s.attributes and s.attributes.get("executor.id") == "executor2" + ] + executor3_processing_spans = [ + s for s in processing_spans if s.attributes and s.attributes.get("executor.id") == "executor3" + ] + + # Get span IDs from processing spans + executor2_processing_span_ids = {format(s.context.span_id, "016x") for s in executor2_processing_spans if s.context} + executor3_processing_span_ids = {format(s.context.span_id, "016x") for s in executor3_processing_spans if s.context} + + executor2_sending_spans = [ + s for s in sending_spans if s.parent and format(s.parent.span_id, "016x") in executor2_processing_span_ids + ] + executor3_sending_spans = [ + s for s in sending_spans if s.parent and format(s.parent.span_id, "016x") in executor3_processing_span_ids + ] + + # Verify that we have sending spans from both executors + assert len(executor2_sending_spans) >= 1, "Should have at least one sending span from executor2" + assert len(executor3_sending_spans) >= 1, "Should have at least one sending span from executor3" + + # Verify that the aggregator span links point to the correct source spans + linked_span_ids = {link.context.span_id for link in aggregator_span.links} + + # Should have links from both executor2 and executor3's sending spans + executor2_span_ids = {s.context.span_id for s in executor2_sending_spans if s.context} + executor3_span_ids = {s.context.span_id for s in executor3_sending_spans if s.context} + + # At least one span from each executor should be linked + assert bool(linked_span_ids & executor2_span_ids), "Aggregator should link to executor2's sending span" + assert bool(linked_span_ids & executor3_span_ids), "Aggregator should link to executor3's sending span" + + # Should have at least 2 links (one from each source executor) + assert len(aggregator_span.links) >= 2, f"Expected at least 2 links, got {len(aggregator_span.links)}" + + +@pytest.mark.asyncio +async def test_workflow_error_handling_in_tracing(tracing_enabled: Any, span_exporter: InMemorySpanExporter) -> None: + """Test that workflow errors are properly recorded in traces.""" + + class FailingExecutor(Executor): + def __init__(self) -> None: + super().__init__(id="failing_executor") + + @handler + async def handle_message(self, message: str, ctx: WorkflowContext[None]) -> None: + raise ValueError("Test error") + + failing_executor = FailingExecutor() + workflow = WorkflowBuilder().set_start_executor(failing_executor).build() + + # Run workflow and expect error + with pytest.raises(ValueError, match="Test error"): + async for _ in workflow.run_streaming("test input"): + pass + + spans = span_exporter.get_finished_spans() + + # Find workflow span + workflow_spans = [s for s in spans if s.name == "workflow.run"] + assert len(workflow_spans) == 1 + + workflow_span = workflow_spans[0] + + # Verify error event and status are recorded + assert workflow_span.events is not None + event_names = [event.name for event in workflow_span.events] + assert "workflow.started" in event_names + assert "workflow.error" in event_names + assert workflow_span.status.status_code.name == "ERROR" + + +@pytest.mark.asyncio +async def test_message_trace_context_serialization() -> None: + """Test that message trace context is properly serialized/deserialized.""" + ctx = InProcRunnerContext() + + # Create message with trace context + message = Message( + data="test", + source_id="source", + target_id="target", + trace_contexts=[{"traceparent": "00-trace-span-01"}], + source_span_ids=["span123"], + ) + + await ctx.send_message(message) + + # Get checkpoint state (which serializes messages) + state = await ctx.get_checkpoint_state() + + # Check serialized message includes trace context + serialized_msg = state["messages"]["source"][0] + assert serialized_msg["trace_contexts"] == [{"traceparent": "00-trace-span-01"}] + assert serialized_msg["source_span_ids"] == ["span123"] + + # Test deserialization + await ctx.set_checkpoint_state(state) + restored_messages = await ctx.drain_messages() + + restored_msg = list(restored_messages.values())[0][0] + assert restored_msg.trace_context == {"traceparent": "00-trace-span-01"} # Test backward compatibility + assert restored_msg.source_span_id == "span123" # Test backward compatibility + assert restored_msg.trace_contexts == [{"traceparent": "00-trace-span-01"}] # Test new format + assert restored_msg.source_span_ids == ["span123"] # Test new format + + +@pytest.mark.asyncio +async def test_workflow_build_error_tracing(tracing_enabled: Any, span_exporter: InMemorySpanExporter) -> None: + """Test that build errors are properly recorded in build spans.""" + + # Test validation error by not setting start executor + builder = WorkflowBuilder() + + with pytest.raises(ValueError, match="Starting executor must be set"): + builder.build() + + spans = span_exporter.get_finished_spans() + assert len(spans) == 1 + + build_span = spans[0] + assert build_span.name == "workflow.build" + + # Verify error status and events + assert build_span.status.status_code.name == "ERROR" + assert build_span.events is not None + + event_names = [event.name for event in build_span.events] + assert "build.started" in event_names + assert "build.error" in event_names + + # Check error event attributes + error_events = [event for event in build_span.events if event.name == "build.error"] + assert len(error_events) == 1 + + error_event = error_events[0] + assert error_event.attributes is not None + assert "Starting executor must be set" in str(error_event.attributes.get("build.error.message")) + assert error_event.attributes.get("build.error.type") == "ValueError" diff --git a/python/packages/workflow/tests/test_workflow.py b/python/packages/workflow/tests/test_workflow.py index 9cf1ad0da4..3d5791ab11 100644 --- a/python/packages/workflow/tests/test_workflow.py +++ b/python/packages/workflow/tests/test_workflow.py @@ -288,7 +288,7 @@ async def test_fan_in(): def simple_executor() -> Executor: class SimpleExecutor(Executor): @handler - async def handle_message(self, message: Message, context: WorkflowContext[None]) -> None: + async def handle_message(self, message: str, context: WorkflowContext[None]) -> None: pass return SimpleExecutor(id="test_executor") diff --git a/python/samples/getting_started/telemetry/.env.example b/python/samples/getting_started/telemetry/.env.example index f4f4322341..5e8bced445 100644 --- a/python/samples/getting_started/telemetry/.env.example +++ b/python/samples/getting_started/telemetry/.env.example @@ -2,3 +2,4 @@ CONNECTION_STRING="..." OTLP_ENDPOINT="http://localhost:4317/" AGENT_FRAMEWORK_GENAI_ENABLE_OTEL_DIAGNOSTICS=true AGENT_FRAMEWORK_GENAI_ENABLE_OTEL_DIAGNOSTICS_SENSITIVE=true +AGENT_FRAMEWORK_WORKFLOW_ENABLE_OTEL_DIAGNOSTICS=true diff --git a/python/samples/getting_started/telemetry/README.md b/python/samples/getting_started/telemetry/README.md index 61bb44fc2a..5d947cda17 100644 --- a/python/samples/getting_started/telemetry/README.md +++ b/python/samples/getting_started/telemetry/README.md @@ -4,9 +4,12 @@ This sample project shows how a Python application can be configured to send Age In this sample, we provide options to send telemetry to [Application Insights](https://learn.microsoft.com/en-us/azure/azure-monitor/app/app-insights-overview), [Aspire Dashboard](https://learn.microsoft.com/en-us/dotnet/aspire/fundamentals/dashboard/overview?tabs=bash), and console output. +> **Quick Start**: For local development without Azure setup, you can use the [Aspire Dashboard](https://learn.microsoft.com/en-us/dotnet/aspire/fundamentals/dashboard/standalone) which runs locally via Docker and provides an excellent telemetry viewing experience for OpenTelemetry data. + > Note that it is also possible to use other Application Performance Management (APM) vendors. An example is [Prometheus](https://prometheus.io/docs/introduction/overview/). Please refer to this [link](https://opentelemetry.io/docs/languages/python/exporters/) to learn more about exporters. For more information, please refer to the following resources: + 1. [Azure Monitor OpenTelemetry Exporter](https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/monitor/azure-monitor-opentelemetry-exporter) 2. [Aspire Dashboard for Python Apps](https://learn.microsoft.com/en-us/dotnet/aspire/fundamentals/dashboard/standalone-for-python?tabs=flask%2Cwindows) 3. [Python Logging](https://docs.python.org/3/library/logging.html) @@ -19,13 +22,18 @@ The Agent Framework Python SDK is designed to efficiently generate comprehensive ## Configuration ### Required resources + 2. OpenAI or [Azure OpenAI](https://learn.microsoft.com/en-us/azure/ai-services/openai/how-to/create-resource?pivots=web-portal) + ### Optional resources + 1. [Application Insights](https://learn.microsoft.com/en-us/azure/azure-monitor/app/create-workspace-resource) 2. [Aspire Dashboard](https://learn.microsoft.com/en-us/dotnet/aspire/fundamentals/dashboard/standalone-for-python?tabs=flask%2Cwindows#start-the-aspire-dashboard) ### Dependencies + You will also need to install the following dependencies to your virtual environment to run this sample: + ```bash # For Azure ApplicationInsights/AzureMonitor uv pip install azure-monitor-opentelemetry azure-monitor-opentelemetry-exporter @@ -39,16 +47,18 @@ uv pip install opentelemetry-exporter-otlp-proto-grpc 2. Create a `.env` file if one doesn't already exist in this folder. Please refer to the [example file](./.env.example). > Note that `CONNECTION_STRING` and `SAMPLE_OTLP_ENDPOINT` are optional. If you don't configure them, everything will get outputted to the console. > Set `AGENT_FRAMEWORK_GENAI_ENABLE_OTEL_DIAGNOSTICS=true` to enable basic telemetry and `AGENT_FRAMEWORK_GENAI_ENABLE_OTEL_DIAGNOSTICS_SENSITIVE=true` to include sensitive information like prompts and responses. + > Set `AGENT_FRAMEWORK_WORKFLOW_ENABLE_OTEL_DIAGNOSTICS=true` to enable workflow telemetry for the workflow samples. > Sensitive information should only be enabled in a development or test environment. It is not recommended to enable this in production environments as it may expose sensitive data. -3. Activate your python virtual environment, and then run `python scenarios.py` or `python interactive.py`. +3. Activate your python virtual environment, and then run `python scenarios.py`, `python interactive.py`, `python agent.py`, or `python workflow.py`. > This will output the Operation/Trace ID, which can be used later for filtering. ### Scenarios -This sample includes two different applications demonstrating Agent Framework telemetry: +This sample includes multiple applications demonstrating Agent Framework telemetry: #### scenarios.py + Organized into specific scenarios where the framework will generate useful telemetry data: - `chat_client`: This is when a chat client is invoked directly (i.e. not streaming) with a weather tool function. **Information about the call to the underlying model and tool usage will be recorded**. @@ -58,8 +68,24 @@ Organized into specific scenarios where the framework will generate useful telem By default, running `python scenarios.py` will run all three scenarios. To run individual scenarios, use the `--scenario` command line argument. For example, `python scenarios.py --scenario chat_client`. For more information, please run `python scenarios.py -h`. #### interactive.py + An interactive chat application that demonstrates telemetry collection in a conversational context. This sample includes the same `get_weather` tool function and allows for multi-turn conversations. Run `python interactive.py` and start chatting. Type 'exit' to quit the application. This sample only logs at the `WARNING` level, so you will not see as much telemetry data as in the `scenarios.py` sample. +#### agent.py + +A sample demonstrating Agent Framework telemetry collection for agent-based workflows. This shows how telemetry is captured when using the Agent Framework's agent abstraction layer, including agent initialization, message processing, and tool execution within an agent context. + +By default, running `python agent.py` will run all agent scenarios. To run individual scenarios, use the `--scenario` command line argument. For example, `python agent.py --scenario basic`. For more information, please run `python agent.py -h`. + +#### workflow.py + +A sample demonstrating workflow telemetry collection for the Agent Framework's workflow execution engine. This includes two scenarios: + +- `sequential`: A simple sequential workflow that processes text through two connected executors (uppercase conversion followed by text reversal). **Information about workflow execution, executor processing, and message passing between executors will be recorded**. +- `sub_workflow`: A more complex scenario demonstrating sub-workflow patterns with a parent workflow orchestrating multiple text processing tasks via sub-workflows. **Information about parent workflow execution, sub-workflow invocation, and cross-workflow communication will be recorded**. + +By default, running `python workflow.py` will run all workflow scenarios. To run individual scenarios, use the `--scenario` command line argument. For example, `python workflow.py --scenario sequential`. For more information, please run `python workflow.py -h`. + ## Application Insights/Azure Monitor ### Logs and traces @@ -100,15 +126,52 @@ dependencies ## Aspire Dashboard +The [Aspire Dashboard](https://learn.microsoft.com/en-us/dotnet/aspire/fundamentals/dashboard/standalone) is a local telemetry viewing tool that provides an excellent experience for viewing OpenTelemetry data without requiring Azure setup. + +### Setting up Aspire Dashboard with Docker + +The easiest way to run the Aspire Dashboard locally is using Docker: + +```bash +# Pull and run the Aspire Dashboard container +docker run --rm -it -d \ + -p 18888:18888 \ + -p 4317:18889 \ + --name aspire-dashboard \ + mcr.microsoft.com/dotnet/aspire-dashboard:latest +``` + +This will start the dashboard with: + +- **Web UI**: Available at +- **OTLP endpoint**: Available at `http://localhost:4317` for your applications to send telemetry data + +### Configuring your application + +Make sure your `.env` file includes the OTLP endpoint: + +```bash +OTLP_ENDPOINT=http://localhost:4317 +``` + +Or set it as an environment variable when running your samples: + +```bash +OTLP_ENDPOINT=http://localhost:4317 python scenarios.py +``` + +### Viewing telemetry data + > Make sure you have the dashboard running to receive telemetry data. -Once the the sample finishes running, navigate to http://localhost:18888 in a web browser to see the telemetry data. Follow the instructions [here](https://learn.microsoft.com/en-us/dotnet/aspire/fundamentals/dashboard/explore) to authenticate to the dashboard and start exploring! +Once your sample finishes running, navigate to in a web browser to see the telemetry data. Follow the [Aspire Dashboard exploration guide](https://learn.microsoft.com/en-us/dotnet/aspire/fundamentals/dashboard/explore) to authenticate to the dashboard and start exploring your traces, logs, and metrics! ## Console output You won't have to deploy an Application Insights resource or install Docker to run Aspire Dashboard if you choose to inspect telemetry data in a console. However, it is difficult to navigate through all the spans and logs produced, so **this method is only recommended when you are just getting started**. We recommend you to get started with the `chat_client` scenario as this generates the least amount of telemetry data. Below is similar to what you will see when you run `python scenarios.py --scenario chat_client`: + ```Json { "name": "chat.completions gpt-4o", diff --git a/python/samples/getting_started/telemetry/workflow.py b/python/samples/getting_started/telemetry/workflow.py new file mode 100644 index 0000000000..d5937c0034 --- /dev/null +++ b/python/samples/getting_started/telemetry/workflow.py @@ -0,0 +1,253 @@ +# Copyright (c) Microsoft. All rights reserved. +# type: ignore +import asyncio +import logging +from typing import Any + +from agent_framework.workflow import ( + Executor, + WorkflowBuilder, + WorkflowCompletedEvent, + WorkflowContext, + handler, +) +from azure.monitor.opentelemetry import configure_azure_monitor +from opentelemetry import trace +from opentelemetry._logs import set_logger_provider +from opentelemetry.exporter.otlp.proto.grpc._log_exporter import OTLPLogExporter +from opentelemetry.exporter.otlp.proto.grpc.metric_exporter import OTLPMetricExporter +from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter +from opentelemetry.metrics import set_meter_provider +from opentelemetry.sdk._logs import LoggerProvider, LoggingHandler +from opentelemetry.sdk._logs.export import BatchLogRecordProcessor, ConsoleLogExporter +from opentelemetry.sdk.metrics import MeterProvider +from opentelemetry.sdk.metrics.export import ConsoleMetricExporter, PeriodicExportingMetricReader +from opentelemetry.sdk.metrics.view import DropAggregation, View +from opentelemetry.sdk.resources import Resource +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import BatchSpanProcessor, ConsoleSpanExporter +from opentelemetry.semconv.attributes import service_attributes +from opentelemetry.trace import SpanKind, set_tracer_provider +from opentelemetry.trace.span import format_trace_id +from pydantic_settings import BaseSettings + +"""Telemetry sample demonstrating OpenTelemetry integration with Agent Framework workflows. + +This sample runs a simple sequential workflow with telemetry collection, +showing telemetry collection for workflow execution, executor processing, +and message publishing between executors. +""" + + +class TelemetrySampleSettings(BaseSettings): + """Settings for the telemetry sample application. + + Optional settings are: + - connection_string: str - The connection string for the Application Insights resource. + This value can be found in the Overview section when examining + your resource from the Azure portal. + (Env var CONNECTION_STRING) + - otlp_endpoint: str - The OTLP endpoint to send telemetry data to. + Depending on the exporter used, you may find this value in different places. + (Env var OTLP_ENDPOINT) + + If no connection string or OTLP endpoint is provided, the telemetry data will be + exported to the console. + """ + + connection_string: str | None = None + otlp_endpoint: str | None = None + + +# Load settings +settings = TelemetrySampleSettings() + +# Create a resource to represent the service/sample +resource = Resource.create({service_attributes.SERVICE_NAME: "WorkflowTelemetryExample"}) + +if settings.connection_string: + configure_azure_monitor( + connection_string=settings.connection_string, + enable_live_metrics=True, + logger_name="agent_framework", + ) + + +def set_up_logging(): + class LogFilter(logging.Filter): + """A filter to not process records from several subpackages.""" + + # These are the namespaces that we want to exclude from logging for the purposes of this demo. + namespaces_to_exclude: list[str] = [ + "httpx", + "openai", + ] + + def filter(self, record): + return not any([record.name.startswith(namespace) for namespace in self.namespaces_to_exclude]) + + exporters = [] + if settings.otlp_endpoint: + exporters.append(OTLPLogExporter(endpoint=settings.otlp_endpoint)) + if not exporters: + exporters.append(ConsoleLogExporter()) + + # Create and set a global logger provider for the application. + logger_provider = LoggerProvider(resource=resource) + # Log processors are initialized with an exporter which is responsible + # for sending the telemetry data to a particular backend. + for log_exporter in exporters: + logger_provider.add_log_record_processor(BatchLogRecordProcessor(log_exporter)) + # Sets the global default logger provider + set_logger_provider(logger_provider) + + # Create a logging handler to write logging records, in OTLP format, to the exporter. + handler = LoggingHandler() + handler.addFilter(LogFilter()) + # Attach the handler to the root logger. `getLogger()` with no arguments returns the root logger. + # Events from all child loggers will be processed by this handler. + logger = logging.getLogger() + logger.addHandler(handler) + # Set the logging level to NOTSET to allow all records to be processed by the handler. + logger.setLevel(logging.NOTSET) + + +def set_up_tracing(): + exporters = [] + if settings.otlp_endpoint: + exporters.append(OTLPSpanExporter(endpoint=settings.otlp_endpoint)) + if not exporters: + exporters.append(ConsoleSpanExporter()) + + # Initialize a trace provider for the application. This is a factory for creating tracers. + tracer_provider = TracerProvider(resource=resource) + # Span processors are initialized with an exporter which is responsible + # for sending the telemetry data to a particular backend. + for exporter in exporters: + tracer_provider.add_span_processor(BatchSpanProcessor(exporter)) + # Sets the global default tracer provider + set_tracer_provider(tracer_provider) + + +def set_up_metrics(): + exporters = [] + if settings.otlp_endpoint: + exporters.append(OTLPMetricExporter(endpoint=settings.otlp_endpoint)) + if not exporters: + exporters.append(ConsoleMetricExporter()) + + # Initialize a metric provider for the application. This is a factory for creating meters. + metric_readers = [ + PeriodicExportingMetricReader(metric_exporter, export_interval_millis=5000) for metric_exporter in exporters + ] + meter_provider = MeterProvider( + metric_readers=metric_readers, + resource=resource, + views=[ + # Dropping all instrument names except for those starting with "agent_framework" + View(instrument_name="*", aggregation=DropAggregation()), + View(instrument_name="agent_framework*"), + ], + ) + # Sets the global default meter provider + set_meter_provider(meter_provider) + + +# Executors for sequential workflow +class UpperCaseExecutor(Executor): + """An executor that converts text to uppercase.""" + + @handler + async def to_upper_case(self, text: str, ctx: WorkflowContext[str]) -> None: + """Execute the task by converting the input string to uppercase.""" + print(f"UpperCaseExecutor: Processing '{text}'") + result = text.upper() + print(f"UpperCaseExecutor: Result '{result}'") + + # Send the result to the next executor in the workflow. + await ctx.send_message(result) + + +class ReverseTextExecutor(Executor): + """An executor that reverses text.""" + + @handler + async def reverse_text(self, text: str, ctx: WorkflowContext[Any]) -> None: + """Execute the task by reversing the input string.""" + print(f"ReverseTextExecutor: Processing '{text}'") + result = text[::-1] + print(f"ReverseTextExecutor: Result '{result}'") + + # Send the result with a workflow completion event. + await ctx.add_event(WorkflowCompletedEvent(result)) + + +async def run_sequential_workflow() -> None: + """Run a simple sequential workflow demonstrating telemetry collection. + + This workflow processes a string through two executors in sequence: + 1. UpperCaseExecutor converts the input to uppercase + 2. ReverseTextExecutor reverses the string and completes the workflow + + Telemetry data collected includes: + - Overall workflow execution spans + - Individual executor processing spans + - Message publishing between executors + - Workflow completion events + """ + + tracer = trace.get_tracer(__name__) + with tracer.start_as_current_span("Scenario: Sequential Workflow", kind=SpanKind.CLIENT) as current_span: + print("Running scenario: Sequential Workflow") + try: + # Step 1: Create the executors. + upper_case_executor = UpperCaseExecutor(id="upper_case_executor") + reverse_text_executor = ReverseTextExecutor(id="reverse_text_executor") + + # Step 2: Build the workflow with the defined edges. + workflow = ( + WorkflowBuilder() + .add_edge(upper_case_executor, reverse_text_executor) + .set_start_executor(upper_case_executor) + .build() + ) + + # Step 3: Run the workflow with an initial message. + input_text = "hello world" + print(f"Starting workflow with input: '{input_text}'") + + completion_event = None + async for event in workflow.run_streaming(input_text): + print(f"Event: {event}") + if isinstance(event, WorkflowCompletedEvent): + # The WorkflowCompletedEvent contains the final result. + completion_event = event + + if completion_event: + print(f"Workflow completed with result: '{completion_event.data}'") + else: + print("Workflow completed without a completion event") + + except Exception as e: + current_span.record_exception(e) + print(f"Error running workflow: {e}") + + +async def main(): + """Run the telemetry sample with a simple sequential workflow.""" + # Set up the providers + # This must be done before any other telemetry calls + set_up_logging() + set_up_tracing() + set_up_metrics() + + tracer = trace.get_tracer("agent_framework") + with tracer.start_as_current_span("Sequential Workflow Scenario", kind=SpanKind.CLIENT) as current_span: + print(f"Trace ID: {format_trace_id(current_span.get_span_context().trace_id)}") + + # Run the sequential workflow scenario + await run_sequential_workflow() + + +if __name__ == "__main__": + asyncio.run(main())