Python: Telemetry and observability follow-up (#833)

* updated telemetry work

* updated telemetry

* slight improvement

* updated tests

* fixes for telemetry

* fixes for mypy

* added settings setup to runner to avoid error

* streamline usage

* updated tests

* updated tests

* further refinement

* fix dumped item for otel

* removed enable_workflow_otel

* final fixes

* final fixes

* updated samples

* removed exporters

* fix tests

* fixed last import'

* fixed devui
This commit is contained in:
Eduard van Valkenburg
2025-09-23 08:21:56 +02:00
committed by GitHub
Unverified
parent f93f16a9ad
commit 2576e7a091
52 changed files with 1625 additions and 1586 deletions
+6 -6
View File
@@ -9,12 +9,6 @@ OPENAI_RESPONSES_MODEL_ID=""
AZURE_OPENAI_ENDPOINT=""
AZURE_OPENAI_CHAT_DEPLOYMENT_NAME=""
AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME=""
# Telemetry
AGENT_FRAMEWORK_MONITOR_CONNECTION_STRING="..."
AGENT_FRAMEWORK_OTLP_ENDPOINT="http://localhost:4317/"
AGENT_FRAMEWORK_ENABLE_OTEL=true
AGENT_FRAMEWORK_ENABLE_SENSITIVE_DATA=true
AGENT_FRAMEWORK_WORKFLOW_ENABLE_OTEL=true
# Mem0
MEM0_API_KEY=""
# Copilot Studio
@@ -25,3 +19,9 @@ COPILOTSTUDIOAGENT__AGENTAPPID=""
# Anthropic
ANTHROPIC_API_KEY=""
ANTHROPIC_MODEL=""
# Observability
ENABLE_OTEL=true
ENABLE_SENSITIVE_DATA=true
OTLP_ENDPOINT="http://localhost:4317/"
# APPLICATIONINSIGHTS_LIVE_METRICS=false
# APPLICATIONINSIGHTS_CONNECTION_STRING="..."
@@ -14,8 +14,8 @@ from agent_framework import (
use_function_invocation,
)
from agent_framework.exceptions import ServiceInitializationError
from agent_framework.observability import use_observability
from agent_framework.openai._chat_client import OpenAIBaseChatClient
from agent_framework.telemetry import use_telemetry
from azure.core.credentials import TokenCredential
from openai.lib.azure import AsyncAzureADTokenProvider, AsyncAzureOpenAI
from openai.types.chat.chat_completion import Choice
@@ -40,7 +40,7 @@ TAzureChatClient = TypeVar("TAzureChatClient", bound="AzureChatClient")
@use_function_invocation
@use_telemetry
@use_observability
class AzureChatClient(AzureOpenAIConfigMixin, OpenAIBaseChatClient):
"""Azure Chat completion class."""
@@ -6,8 +6,8 @@ from urllib.parse import urljoin
from agent_framework import use_function_invocation
from agent_framework.exceptions import ServiceInitializationError
from agent_framework.observability import use_observability
from agent_framework.openai._responses_client import OpenAIBaseResponsesClient
from agent_framework.telemetry import use_telemetry
from azure.core.credentials import TokenCredential
from openai.lib.azure import AsyncAzureADTokenProvider, AsyncAzureOpenAI
from pydantic import SecretStr, ValidationError
@@ -21,7 +21,7 @@ from ._shared import (
TAzureResponsesClient = TypeVar("TAzureResponsesClient", bound="AzureResponsesClient")
@use_telemetry
@use_observability
@use_function_invocation
class AzureResponsesClient(AzureOpenAIConfigMixin, OpenAIBaseResponsesClient):
"""Azure Responses completion class."""
@@ -7,9 +7,9 @@ from copy import copy
from typing import Any, ClassVar, Final
from agent_framework._pydantic import AFBaseSettings, HTTPsUrl
from agent_framework._telemetry import APP_INFO, USER_AGENT_KEY, prepend_agent_framework_to_user_agent
from agent_framework.exceptions import ServiceInitializationError
from agent_framework.openai._shared import OpenAIBase
from agent_framework.telemetry import APP_INFO, USER_AGENT_KEY, prepend_agent_framework_to_user_agent
from azure.core.credentials import TokenCredential
from openai.lib.azure import AsyncAzureOpenAI
from pydantic import ConfigDict, SecretStr, model_validator, validate_call
@@ -18,12 +18,12 @@ from agent_framework import (
TextContent,
ai_function,
)
from agent_framework._telemetry import USER_AGENT_KEY
from agent_framework.exceptions import ServiceInitializationError, ServiceResponseException
from agent_framework.openai import (
ContentFilterResultSeverity,
OpenAIContentFilterException,
)
from agent_framework.telemetry import USER_AGENT_KEY
from azure.identity import AzureCliCredential
from httpx import Request, Response
from openai import AsyncAzureOpenAI, AsyncStream
@@ -72,17 +72,17 @@ class AgentFrameworkExecutor:
def _setup_agent_framework_tracing(self) -> None:
"""Set up Agent Framework's built-in tracing."""
# Configure Agent Framework tracing only if OTLP endpoint is configured
otlp_endpoint = os.environ.get("AGENT_FRAMEWORK_OTLP_ENDPOINT")
otlp_endpoint = os.environ.get("OTLP_ENDPOINT")
if otlp_endpoint:
try:
from agent_framework.telemetry import setup_telemetry
from agent_framework.observability import setup_observability
setup_telemetry(enable_otel=True, enable_sensitive_data=True, otlp_endpoint=otlp_endpoint)
logger.info(f"Enabled Agent Framework telemetry with endpoint: {otlp_endpoint}")
setup_observability(enable_sensitive_data=True, otlp_endpoint=otlp_endpoint)
logger.info(f"Enabled Agent Framework observability with endpoint: {otlp_endpoint}")
except Exception as e:
logger.warning(f"Failed to enable Agent Framework tracing: {e}")
logger.warning(f"Failed to enable Agent Framework observability: {e}")
else:
logger.debug("No OTLP endpoint configured, skipping telemetry setup")
logger.debug("No OTLP endpoint configured, skipping observability setup")
# Thread Management Methods
def create_thread(self, agent_id: str) -> str:
@@ -6,6 +6,7 @@ from collections.abc import AsyncIterable, MutableMapping, MutableSequence
from typing import Any, ClassVar, TypeVar
from agent_framework import (
AGENT_FRAMEWORK_USER_AGENT,
AIFunction,
BaseChatClient,
ChatMessage,
@@ -27,7 +28,7 @@ from agent_framework import (
)
from agent_framework._pydantic import AFBaseSettings
from agent_framework.exceptions import ServiceInitializationError, ServiceResponseException
from agent_framework.telemetry import AGENT_FRAMEWORK_USER_AGENT, use_telemetry
from agent_framework.observability import use_observability
from azure.ai.agents.models import (
AgentsNamedToolChoice,
AgentsNamedToolChoiceType,
@@ -97,7 +98,7 @@ TFoundryChatClient = TypeVar("TFoundryChatClient", bound="FoundryChatClient")
@use_function_invocation
@use_telemetry
@use_observability
class FoundryChatClient(BaseChatClient):
"""Azure AI Foundry Chat client."""
@@ -190,17 +191,17 @@ class FoundryChatClient(BaseChatClient):
)
self._should_close_client = should_close_client
async def setup_foundry_telemetry(self, enable_live_metrics: bool = False) -> None:
async def setup_foundry_observability(self, enable_live_metrics: bool = False) -> None:
"""Call this method to setup tracing with Foundry.
This will take the connection string from the project client.
It will override any connection string that is set in the environment variables.
It will disable any OTLP endpoint that might have been set.
"""
from agent_framework.telemetry import setup_telemetry
from agent_framework.observability import setup_observability
setup_telemetry(
application_insights_connection_string=await self.client.telemetry.get_application_insights_connection_string(), # noqa: E501
setup_observability(
applicationinsights_connection_string=await self.client.telemetry.get_application_insights_connection_string(), # noqa: E501
enable_live_metrics=enable_live_metrics,
)
@@ -32,7 +32,7 @@ class GAIATelemetryConfig:
self,
enable_tracing: bool = False,
otlp_endpoint: str | None = None,
application_insights_connection_string: str | None = None,
applicationinsights_connection_string: str | None = None,
enable_live_metrics: bool = False,
trace_to_file: bool = False,
file_path: str | None = None,
@@ -43,30 +43,29 @@ class GAIATelemetryConfig:
Args:
enable_tracing: Whether to enable OpenTelemetry tracing
otlp_endpoint: OTLP endpoint for trace export
application_insights_connection_string: Azure Monitor connection string
applicationinsights_connection_string: Azure Monitor connection string
enable_live_metrics: Enable Azure Monitor live metrics
trace_to_file: Whether to export traces to local file
file_path: Path for local file export (defaults to gaia_traces.json)
"""
self.enable_tracing = enable_tracing
self.otlp_endpoint = otlp_endpoint
self.application_insights_connection_string = application_insights_connection_string
self.applicationinsights_connection_string = applicationinsights_connection_string
self.enable_live_metrics = enable_live_metrics
self.trace_to_file = trace_to_file
self.file_path = file_path or "gaia_traces.json"
def setup_telemetry(self) -> None:
def setup_observability(self) -> None:
"""Set up OpenTelemetry based on configuration."""
if not self.enable_tracing:
return
from agent_framework.telemetry import setup_telemetry
from agent_framework.observability import setup_observability
setup_telemetry(
enable_otel=True,
setup_observability(
enable_sensitive_data=True, # Enable for detailed task traces
otlp_endpoint=self.otlp_endpoint,
application_insights_connection_string=self.application_insights_connection_string,
applicationinsights_connection_string=self.applicationinsights_connection_string,
enable_live_metrics=self.enable_live_metrics,
)
@@ -272,7 +271,7 @@ class GAIA:
self.telemetry_config = telemetry_config or GAIATelemetryConfig()
# Set up telemetry
self.telemetry_config.setup_telemetry()
self.telemetry_config.setup_observability()
# Initialize tracer
if self.telemetry_config.enable_tracing:
@@ -29,7 +29,7 @@ async def main() -> None:
enable_tracing=True, # Enable OpenTelemetry tracing
# Optional: Configure external endpoints
# otlp_endpoint="http://localhost:4317", # For Aspire Dashboard or other OTLP endpoints
# application_insights_connection_string="your_connection_string", # For Azure Monitor
# applicationinsights_connection_string="your_connection_string", # For Azure Monitor
# enable_live_metrics=True, # Enable Azure Monitor live metrics
# Configure local file tracing
trace_to_file=True, # Export traces to local file
@@ -2,11 +2,13 @@
import importlib
import importlib.metadata
from typing import Final
try:
__version__ = importlib.metadata.version(__name__)
_version = importlib.metadata.version(__name__)
except importlib.metadata.PackageNotFoundError:
__version__ = "0.0.0" # Fallback for development mode
_version = "0.0.0" # Fallback for development mode
__version__: Final[str] = _version
from ._agents import * # noqa: F403
from ._clients import * # noqa: F403
@@ -14,6 +16,7 @@ from ._logging import * # noqa: F403
from ._mcp import * # noqa: F403
from ._memory import * # noqa: F403
from ._middleware import * # noqa: F403
from ._telemetry import * # noqa: F403
from ._threads import * # noqa: F403
from ._tools import * # noqa: F403
from ._types import * # noqa: F403
@@ -28,7 +28,7 @@ from ._types import (
Role,
)
from .exceptions import AgentExecutionException
from .telemetry import use_agent_telemetry
from .observability import use_agent_observability
if sys.version_info >= (3, 11):
from typing import Self # pragma: no cover
@@ -258,7 +258,7 @@ class BaseAgent(AFBaseModel):
@use_agent_middleware
@use_agent_telemetry
@use_agent_observability
class ChatAgent(BaseAgent):
"""A Chat Client Agent."""
@@ -428,7 +428,6 @@ class BaseChatClient(AFBaseModel, ABC):
tools=self._normalize_tools(tools), # type: ignore
user=user,
additional_properties=additional_properties or {},
**kwargs,
)
prepped_messages = self.prepare_messages(messages)
self._prepare_tool_choice(chat_options=chat_options)
@@ -0,0 +1,59 @@
# Copyright (c) Microsoft. All rights reserved.
import os
from typing import Any, Final
from . import __version__ as version_info
from ._logging import get_logger
logger = get_logger()
__all__ = [
"AGENT_FRAMEWORK_USER_AGENT",
"APP_INFO",
"USER_AGENT_KEY",
"USER_AGENT_TELEMETRY_DISABLED_ENV_VAR",
"prepend_agent_framework_to_user_agent",
]
# Note that if this environment variable does not exist, user agent telemetry is enabled.
USER_AGENT_TELEMETRY_DISABLED_ENV_VAR = "AGENT_FRAMEWORK_USER_AGENT_DISABLED"
IS_TELEMETRY_ENABLED = os.environ.get(USER_AGENT_TELEMETRY_DISABLED_ENV_VAR, "false").lower() not in ["true", "1"]
APP_INFO = (
{
"agent-framework-version": f"python/{version_info}", # type: ignore[has-type]
}
if IS_TELEMETRY_ENABLED
else None
)
USER_AGENT_KEY: Final[str] = "User-Agent"
HTTP_USER_AGENT: Final[str] = "agent-framework-python"
AGENT_FRAMEWORK_USER_AGENT = f"{HTTP_USER_AGENT}/{version_info}" # type: ignore[has-type]
def prepend_agent_framework_to_user_agent(headers: dict[str, Any] | None = None) -> dict[str, Any]:
"""Prepend "agent-framework" to the User-Agent in the headers.
When user agent telemetry is disabled, through the AZURE_TELEMETRY_DISABLED environment variable,
the User-Agent header will not include the agent-framework information, it will be sent back as is,
or as a empty dict when None is passed.
Args:
headers: The existing headers dictionary.
Returns:
A new dict with "User-Agent" set to "agent-framework-python/{version}" if headers is None.
The modified headers dictionary with "agent-framework-python/{version}" prepended to the User-Agent.
"""
if not IS_TELEMETRY_ENABLED:
return headers or {}
if not headers:
return {USER_AGENT_KEY: AGENT_FRAMEWORK_USER_AGENT}
headers[USER_AGENT_KEY] = (
f"{AGENT_FRAMEWORK_USER_AGENT} {headers[USER_AGENT_KEY]}"
if USER_AGENT_KEY in headers
else AGENT_FRAMEWORK_USER_AGENT
)
return headers
+17 -15
View File
@@ -21,19 +21,19 @@ from typing import (
runtime_checkable,
)
from opentelemetry import metrics
from opentelemetry.metrics import Histogram
from pydantic import AnyUrl, BaseModel, Field, PrivateAttr, ValidationError, create_model, field_validator
from ._logging import get_logger
from ._pydantic import AFBaseModel
from .exceptions import ChatClientInitializationError, ToolException
from .telemetry import (
from .observability import (
OPERATION_DURATION_BUCKET_BOUNDARIES,
OtelAttr,
_capture_exception, # type: ignore
capture_exception, # type: ignore
get_function_span,
get_function_span_attributes,
meter,
get_meter,
)
if TYPE_CHECKING:
@@ -358,6 +358,16 @@ class HostedFileSearchTool(BaseTool):
super().__init__(**args, **kwargs)
def _default_histogram() -> Histogram:
"""Get the default histogram for function invocation duration."""
return get_meter().create_histogram(
name=OtelAttr.MEASUREMENT_FUNCTION_INVOCATION_DURATION,
unit=OtelAttr.DURATION_UNIT,
description="Measures the duration of a function's execution",
explicit_bucket_boundaries_advisory=OPERATION_DURATION_BUCKET_BOUNDARIES,
)
class AIFunction(BaseTool, Generic[ArgsT, ReturnT]):
"""A AITool that is callable as code.
@@ -371,14 +381,7 @@ class AIFunction(BaseTool, Generic[ArgsT, ReturnT]):
func: Callable[..., Awaitable[ReturnT] | ReturnT]
input_model: type[ArgsT]
_invocation_duration_histogram: metrics.Histogram = PrivateAttr(
default_factory=lambda: meter.create_histogram(
name=OtelAttr.MEASUREMENT_FUNCTION_INVOCATION_DURATION,
unit=OtelAttr.DURATION_UNIT,
description="Measures the duration of a function's execution",
explicit_bucket_boundaries_advisory=OPERATION_DURATION_BUCKET_BOUNDARIES,
)
)
_invocation_duration_histogram: Histogram = PrivateAttr(default_factory=_default_histogram)
def __call__(self, *args: Any, **kwargs: Any) -> ReturnT | Awaitable[ReturnT]:
"""Call the wrapped function with the provided arguments."""
@@ -398,7 +401,7 @@ class AIFunction(BaseTool, Generic[ArgsT, ReturnT]):
kwargs: keyword arguments to pass to the function, will not be used if `arguments` is provided.
"""
global OTEL_SETTINGS
from .telemetry import OTEL_SETTINGS, setup_telemetry
from .observability import OTEL_SETTINGS
tool_call_id = kwargs.pop("tool_call_id", None)
if arguments is not None:
@@ -414,7 +417,6 @@ class AIFunction(BaseTool, Generic[ArgsT, ReturnT]):
logger.debug(f"Function result: {result or 'None'}")
return result # type: ignore[reportReturnType]
setup_telemetry()
attributes = get_function_span_attributes(self, tool_call_id=tool_call_id)
if OTEL_SETTINGS.SENSITIVE_DATA_ENABLED: # type: ignore[name-defined]
attributes.update({
@@ -438,7 +440,7 @@ class AIFunction(BaseTool, Generic[ArgsT, ReturnT]):
except Exception as exception:
end_time_stamp = perf_counter()
attributes[OtelAttr.ERROR_TYPE] = type(exception).__name__
_capture_exception(span=span, exception=exception, timestamp=time_ns())
capture_exception(span=span, exception=exception, timestamp=time_ns())
logger.error(f"Function failed. Error: {exception}")
raise
else:
@@ -88,7 +88,6 @@ from ._runner_context import (
)
from ._sequential import SequentialBuilder
from ._shared_state import SharedState
from ._telemetry import EdgeGroupDeliveryStatus, WorkflowTracer, workflow_tracer
from ._validation import (
EdgeDuplicationError,
ExecutorDuplicationError,
@@ -116,7 +115,6 @@ __all__ = [
"Default",
"Edge",
"EdgeDuplicationError",
"EdgeGroupDeliveryStatus",
"Executor",
"ExecutorCompletedEvent",
"ExecutorDuplicationError",
@@ -184,7 +182,6 @@ __all__ = [
"WorkflowRunState",
"WorkflowStartedEvent",
"WorkflowStatusEvent",
"WorkflowTracer",
"WorkflowValidationError",
"WorkflowViz",
"create_edge_runner",
@@ -192,7 +189,6 @@ __all__ = [
"handler",
"intercepts_request",
"validate_workflow_graph",
"workflow_tracer",
]
# Rebuild models to resolve forward references after all imports are complete
@@ -84,7 +84,6 @@ from ._runner_context import (
)
from ._sequential import SequentialBuilder
from ._shared_state import SharedState
from ._telemetry import EdgeGroupDeliveryStatus, WorkflowTracer, workflow_tracer
from ._validation import (
EdgeDuplicationError,
ExecutorDuplicationError,
@@ -112,7 +111,6 @@ __all__ = [
"Default",
"Edge",
"EdgeDuplicationError",
"EdgeGroupDeliveryStatus",
"Executor",
"ExecutorCompletedEvent",
"ExecutorDuplicationError",
@@ -180,7 +178,6 @@ __all__ = [
"WorkflowRunState",
"WorkflowStartedEvent",
"WorkflowStatusEvent",
"WorkflowTracer",
"WorkflowValidationError",
"WorkflowViz",
"create_edge_runner",
@@ -188,5 +185,4 @@ __all__ = [
"handler",
"intercepts_request",
"validate_workflow_graph",
"workflow_tracer",
]
@@ -20,9 +20,9 @@ from agent_framework import (
TextContent,
UsageDetails,
)
from agent_framework._pydantic import AFBaseModel
from agent_framework.exceptions import AgentExecutionException
from .._pydantic import AFBaseModel
from ..exceptions import AgentExecutionException
from ._events import (
AgentRunUpdateEvent,
RequestInfoEvent,
@@ -8,8 +8,7 @@ from typing import Any, ClassVar
from pydantic import Field
from agent_framework._pydantic import AFBaseModel
from .._pydantic import AFBaseModel
from ._executor import Executor
logger = logging.getLogger(__name__)
@@ -6,11 +6,11 @@ from abc import ABC, abstractmethod
from collections import defaultdict
from typing import Any
from ..observability import EdgeGroupDeliveryStatus, OtelAttr, create_edge_group_processing_span
from ._edge import Edge, EdgeGroup, FanInEdgeGroup, FanOutEdgeGroup, SingleEdgeGroup, SwitchCaseEdgeGroup
from ._executor import Executor
from ._runner_context import Message, RunnerContext
from ._shared_state import SharedState
from ._telemetry import EdgeGroupDeliveryStatus, workflow_tracer
from ._workflow_context import WorkflowContext
logger = logging.getLogger(__name__)
@@ -90,39 +90,49 @@ class SingleEdgeRunner(EdgeRunner):
should_execute = False
target_id = None
source_id = None
with workflow_tracer.create_edge_group_processing_span(
with create_edge_group_processing_span(
self._edge_group.__class__.__name__,
edge_group_id=self._edge_group.id,
message_source_id=message.source_id,
message_target_id=message.target_id,
source_trace_contexts=message.trace_contexts,
source_span_ids=message.source_span_ids,
):
) as span:
try:
if message.target_id and message.target_id != self._edge.target_id:
workflow_tracer.set_edge_group_span_attributes(
False, EdgeGroupDeliveryStatus.DROPPED_TARGET_MISMATCH
)
span.set_attributes({
OtelAttr.EDGE_GROUP_DELIVERED: False,
OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DROPPED_TARGET_MISMATCH.value,
})
return False
if self._can_handle(self._edge.target_id, message.data):
if self._edge.should_route(message.data):
workflow_tracer.set_edge_group_span_attributes(True, EdgeGroupDeliveryStatus.DELIVERED)
span.set_attributes({
OtelAttr.EDGE_GROUP_DELIVERED: True,
OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DELIVERED.value,
})
should_execute = True
target_id = self._edge.target_id
source_id = self._edge.source_id
else:
workflow_tracer.set_edge_group_span_attributes(
False, EdgeGroupDeliveryStatus.DROPPED_CONDITION_FALSE
)
span.set_attributes({
OtelAttr.EDGE_GROUP_DELIVERED: False,
OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DROPPED_CONDITION_FALSE.value,
})
# Return True here because message was processed, just condition failed
return True
else:
workflow_tracer.set_edge_group_span_attributes(False, EdgeGroupDeliveryStatus.DROPPED_TYPE_MISMATCH)
span.set_attributes({
OtelAttr.EDGE_GROUP_DELIVERED: False,
OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DROPPED_TYPE_MISMATCH.value,
})
return False
except Exception as e:
workflow_tracer.set_edge_group_span_attributes(False, EdgeGroupDeliveryStatus.EXCEPTION)
span.set_attributes({
OtelAttr.EDGE_GROUP_DELIVERED: False,
OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.EXCEPTION.value,
})
raise e
# Execute outside the span
@@ -147,22 +157,24 @@ class FanOutEdgeRunner(EdgeRunner):
"""Send a message through all edges in the fan-out edge group."""
deliverable_edges = []
single_target_edge = None
# Process routing logic within span
with workflow_tracer.create_edge_group_processing_span(
with create_edge_group_processing_span(
self._edge_group.__class__.__name__,
edge_group_id=self._edge_group.id,
message_source_id=message.source_id,
message_target_id=message.target_id,
source_trace_contexts=message.trace_contexts,
source_span_ids=message.source_span_ids,
):
) as span:
try:
selection_results = (
self._selection_func(message.data, self._target_ids) if self._selection_func else self._target_ids
)
if not self._validate_selection_result(selection_results):
workflow_tracer.set_edge_group_span_attributes(False, EdgeGroupDeliveryStatus.EXCEPTION)
span.set_attributes({
OtelAttr.EDGE_GROUP_DELIVERED: False,
OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.EXCEPTION.value,
})
raise RuntimeError(
f"Invalid selection result: {selection_results}. "
f"Expected selections to be a subset of valid target executor IDs: {self._target_ids}."
@@ -174,24 +186,30 @@ 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):
workflow_tracer.set_edge_group_span_attributes(True, EdgeGroupDeliveryStatus.DELIVERED)
span.set_attributes({
OtelAttr.EDGE_GROUP_DELIVERED: True,
OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DELIVERED.value,
})
single_target_edge = edge
else:
workflow_tracer.set_edge_group_span_attributes(
False, EdgeGroupDeliveryStatus.DROPPED_CONDITION_FALSE
)
span.set_attributes({
OtelAttr.EDGE_GROUP_DELIVERED: False,
OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DROPPED_CONDITION_FALSE.value, # noqa: E501
})
# For targeted messages with condition failure, return True (message was processed)
return True
else:
workflow_tracer.set_edge_group_span_attributes(
False, EdgeGroupDeliveryStatus.DROPPED_TYPE_MISMATCH
)
span.set_attributes({
OtelAttr.EDGE_GROUP_DELIVERED: False,
OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DROPPED_TYPE_MISMATCH.value, # noqa: E501
})
# For targeted messages that can't be handled, return False
return False
else:
workflow_tracer.set_edge_group_span_attributes(
False, EdgeGroupDeliveryStatus.DROPPED_TARGET_MISMATCH
)
span.set_attributes({
OtelAttr.EDGE_GROUP_DELIVERED: False,
OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DROPPED_TARGET_MISMATCH.value,
})
# For targeted messages not in selection, return False
return False
else:
@@ -202,14 +220,21 @@ class FanOutEdgeRunner(EdgeRunner):
deliverable_edges.append(edge)
if len(deliverable_edges) > 0:
workflow_tracer.set_edge_group_span_attributes(True, EdgeGroupDeliveryStatus.DELIVERED)
span.set_attributes({
OtelAttr.EDGE_GROUP_DELIVERED: True,
OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DELIVERED.value,
})
else:
workflow_tracer.set_edge_group_span_attributes(
False, EdgeGroupDeliveryStatus.DROPPED_TYPE_MISMATCH
)
span.set_attributes({
OtelAttr.EDGE_GROUP_DELIVERED: False,
OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DROPPED_TYPE_MISMATCH.value,
})
except Exception as e:
workflow_tracer.set_edge_group_span_attributes(False, EdgeGroupDeliveryStatus.EXCEPTION)
span.set_attributes({
OtelAttr.EDGE_GROUP_DELIVERED: False,
OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.EXCEPTION.value,
})
raise e
# Execute outside the span
@@ -250,30 +275,36 @@ class FanInEdgeRunner(EdgeRunner):
async def send_message(self, message: Message, shared_state: SharedState, ctx: RunnerContext) -> bool:
"""Send a message through all edges in the fan-in edge group."""
execution_data: dict[str, Any] | None = None
with workflow_tracer.create_edge_group_processing_span(
with create_edge_group_processing_span(
self._edge_group.__class__.__name__,
edge_group_id=self._edge_group.id,
message_source_id=message.source_id,
message_target_id=message.target_id,
source_trace_contexts=message.trace_contexts,
source_span_ids=message.source_span_ids,
):
) as span:
try:
if message.target_id and message.target_id != self._edges[0].target_id:
workflow_tracer.set_edge_group_span_attributes(
False, EdgeGroupDeliveryStatus.DROPPED_TARGET_MISMATCH
)
span.set_attributes({
OtelAttr.EDGE_GROUP_DELIVERED: False,
OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DROPPED_TARGET_MISMATCH.value,
})
return False
# Check if target can handle list of message data (fan-in aggregates multiple messages)
if self._can_handle(self._edges[0].target_id, [message.data]):
# If the edge can handle the data, buffer the message
self._buffer[message.source_id].append(message)
workflow_tracer.set_edge_group_span_attributes(True, EdgeGroupDeliveryStatus.BUFFERED)
span.set_attributes({
OtelAttr.EDGE_GROUP_DELIVERED: True,
OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.BUFFERED.value,
})
else:
# If the edge cannot handle the data, return False
workflow_tracer.set_edge_group_span_attributes(False, EdgeGroupDeliveryStatus.DROPPED_TYPE_MISMATCH)
span.set_attributes({
OtelAttr.EDGE_GROUP_DELIVERED: False,
OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DROPPED_TYPE_MISMATCH.value,
})
return False
if self._is_ready_to_send():
@@ -294,7 +325,10 @@ class FanInEdgeRunner(EdgeRunner):
trace_contexts=trace_contexts,
source_span_ids=source_span_ids,
)
workflow_tracer.set_edge_group_span_attributes(True, EdgeGroupDeliveryStatus.DELIVERED)
span.set_attributes({
OtelAttr.EDGE_GROUP_DELIVERED: True,
OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DELIVERED.value,
})
# Store execution data for later
execution_data = {
@@ -304,7 +338,10 @@ class FanInEdgeRunner(EdgeRunner):
}
except Exception as e:
workflow_tracer.set_edge_group_span_attributes(False, EdgeGroupDeliveryStatus.EXCEPTION)
span.set_attributes({
OtelAttr.EDGE_GROUP_DELIVERED: False,
OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.EXCEPTION.value,
})
raise e
# Execute outside the span if needed
@@ -12,14 +12,13 @@ from textwrap import shorten
from types import UnionType
from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar, Union, cast, get_args, get_origin, overload
if TYPE_CHECKING:
from ._workflow import Workflow
from pydantic import Field
from agent_framework import AgentProtocol, AgentRunResponse, AgentRunResponseUpdate, AgentThread, ChatMessage
from agent_framework._pydantic import AFBaseModel
from .._agents import AgentProtocol
from .._pydantic import AFBaseModel
from .._threads import AgentThread
from .._types import AgentRunResponse, AgentRunResponseUpdate, ChatMessage
from ..observability import create_processing_span
from ._checkpoint import WorkflowCheckpoint
from ._events import (
AgentRunEvent,
@@ -27,14 +26,16 @@ from ._events import (
ExecutorCompletedEvent,
ExecutorInvokedEvent,
RequestInfoEvent,
_framework_event_origin, # pyright: ignore[reportPrivateUsage]
_framework_event_origin, # type: ignore[reportPrivateUsage]
)
from ._runner_context import _decode_checkpoint_value
from ._runner_context import _decode_checkpoint_value # type: ignore[reportPrivateUsage]
from ._typing_utils import is_instance_of
from ._workflow_context import WorkflowContext
logger = logging.getLogger(__name__)
if TYPE_CHECKING:
from ._workflow import Workflow
logger = logging.getLogger(__name__)
# region Executor
@@ -114,7 +115,6 @@ class Executor(AFBaseModel):
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)
@@ -125,7 +125,7 @@ class Executor(AFBaseModel):
if isinstance(message, Message):
message = message.data
with workflow_tracer.create_processing_span(
with create_processing_span(
self.id,
self.__class__.__name__,
type(message).__name__,
@@ -1,310 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
from enum import Enum
from typing import TYPE_CHECKING, Any, ClassVar
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
from agent_framework._pydantic import AFBaseSettings
if TYPE_CHECKING:
from ._workflow import Workflow
class EdgeGroupDeliveryStatus(Enum):
"""Enum for edge group delivery status values."""
DELIVERED = "delivered"
DROPPED_TYPE_MISMATCH = "dropped type mismatch"
DROPPED_TARGET_MISMATCH = "dropped target mismatch"
DROPPED_CONDITION_FALSE = "dropped condition evaluated to false"
EXCEPTION = "exception"
BUFFERED = "buffered"
# Span name constants
_WORKFLOW_BUILD_SPAN = "workflow.build"
_WORKFLOW_RUN_SPAN = "workflow.run"
_EXECUTOR_PROCESS_SPAN = "executor.process"
_MESSAGE_SEND_SPAN = "message.send"
_EDGE_GROUP_PROCESS_SPAN = "edge_group.process"
class WorkflowDiagnosticSettings(AFBaseSettings):
"""Settings for workflow tracing diagnostics."""
env_prefix: ClassVar[str] = "AGENT_FRAMEWORK_WORKFLOW_"
enable_otel: bool = False
@property
def ENABLED(self) -> bool:
return self.enable_otel
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)
- Edge group processing spans (edge_group.process)
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 create_edge_group_processing_span(
self,
edge_group_type: str,
edge_group_id: str | None = None,
message_source_id: str | None = None,
message_target_id: str | None = None,
source_trace_contexts: list[dict[str, str]] | None = None,
source_span_ids: list[str] | None = None,
) -> Any:
"""Create an edge group processing span with optional links to source spans.
Edge group processing spans track the processing operations in edge runners
before message delivery, including condition checking and routing decisions.
Links to source spans provide causality tracking without unwanted nesting.
Args:
edge_group_type: The type of the edge group (class name).
edge_group_id: The unique ID of the edge group.
message_source_id: The source ID of the message being processed.
message_target_id: The target ID of the message being processed.
source_trace_contexts: Optional trace contexts from source spans for linking.
source_span_ids: Optional source span IDs for linking.
"""
attributes: dict[str, str] = {
"edge_group.type": edge_group_type,
}
if edge_group_id is not None:
attributes["edge_group.id"] = edge_group_id
if message_source_id is not None:
attributes["message.source_id"] = message_source_id
if message_target_id is not None:
attributes["message.target_id"] = message_target_id
# 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(
_EDGE_GROUP_PROCESS_SPAN,
kind=SpanKind.INTERNAL,
attributes=attributes,
links=links,
)
def set_edge_group_span_attributes(self, delivered: bool, delivery_status: EdgeGroupDeliveryStatus) -> None:
"""Set edge group span attributes for delivery status.
Args:
delivered: Whether the message was delivered.
delivery_status: The delivery status from EdgeGroupDeliveryStatus enum.
"""
span = get_current_span()
if span and span.is_recording():
span.set_attributes({
"edge_group.delivered": delivered,
"edge_group.delivery_status": delivery_status.value,
})
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(by_alias=True),
})
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()
@@ -11,9 +11,9 @@ from typing import Any
from pydantic import Field
from agent_framework import AgentProtocol
from agent_framework._pydantic import AFBaseModel
from .._agents import AgentProtocol
from .._pydantic import AFBaseModel
from ..observability import OtelAttr, capture_exception, create_workflow_span
from ._agent import WorkflowAgent
from ._checkpoint import CheckpointStorage
from ._const import DEFAULT_MAX_ITERATIONS
@@ -242,17 +242,19 @@ class Workflow(AFBaseModel):
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):
with create_workflow_span(
OtelAttr.WORKFLOW_RUN_SPAN,
{
OtelAttr.WORKFLOW_ID: self.id,
},
) as span:
saw_completed = False
saw_request = False
emitted_in_progress_pending = False
try:
# Add workflow started event (telemetry + surface state to consumers)
workflow_tracer.add_workflow_event("workflow.started")
span.add_event(OtelAttr.WORKFLOW_STARTED)
# Emit explicit start/status events to the stream
with _framework_event_origin():
started = WorkflowStartedEvent()
@@ -298,17 +300,24 @@ class Workflow(AFBaseModel):
terminal_status = WorkflowStatusEvent(WorkflowRunState.IDLE)
yield terminal_status
workflow_tracer.add_workflow_event("workflow.completed")
except Exception as e:
span.add_event(OtelAttr.WORKFLOW_COMPLETED)
except Exception as exc:
# Surface structured failure details before propagating exception
details = WorkflowErrorDetails.from_exception(e)
details = WorkflowErrorDetails.from_exception(exc)
with _framework_event_origin():
failed_event = WorkflowFailedEvent(details)
yield failed_event
with _framework_event_origin():
failed_status = WorkflowStatusEvent(WorkflowRunState.FAILED)
yield failed_status
workflow_tracer.add_workflow_error_event(e)
span.add_event(
name=OtelAttr.WORKFLOW_ERROR,
attributes={
"error.message": str(exc),
"error.type": type(exc).__name__,
},
)
capture_exception(span, exception=exc)
raise
async def run_stream(self, message: Any) -> AsyncIterable[WorkflowEvent]:
@@ -1074,14 +1083,11 @@ class WorkflowBuilder:
WorkflowValidationError: If workflow validation fails (includes EdgeDuplicationError,
TypeCompatibilityError, and GraphConnectivityError subclasses).
"""
# Import here to avoid circular imports
from ._telemetry import workflow_tracer
# Create workflow build span that includes validation and workflow creation
with workflow_tracer.create_workflow_build_span():
with create_workflow_span(OtelAttr.WORKFLOW_BUILD_SPAN) as span:
try:
# Add workflow build started event
workflow_tracer.add_build_event("build.started")
span.add_event(OtelAttr.BUILD_STARTED)
if not self._start_executor:
raise ValueError(
@@ -1097,7 +1103,7 @@ class WorkflowBuilder:
)
# Add validation completed event
workflow_tracer.add_build_event("build.validation_completed")
span.add_event(OtelAttr.BUILD_VALIDATION_COMPLETED)
context = InProcRunnerContext(self._checkpoint_storage)
@@ -1105,16 +1111,21 @@ class WorkflowBuilder:
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)
span.set_attributes({
OtelAttr.WORKFLOW_ID: workflow.id,
OtelAttr.WORKFLOW_DEFINITION: workflow.model_dump_json(by_alias=True),
})
# Add workflow build completed event
workflow_tracer.add_build_event("build.completed")
span.add_event(OtelAttr.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)
except Exception as exc:
attributes = {
OtelAttr.BUILD_ERROR_MESSAGE: str(exc),
OtelAttr.BUILD_ERROR_TYPE: type(exc).__name__,
}
span.add_event(OtelAttr.BUILD_ERROR, attributes) # type: ignore[reportArgumentType, arg-type]
capture_exception(span, exc)
raise
@@ -4,7 +4,9 @@ import logging
from typing import Any, Generic, TypeVar, cast, get_args
from opentelemetry.propagate import inject
from opentelemetry.trace import SpanKind
from ..observability import OtelAttr, create_workflow_span
from ._events import (
WorkflowEvent,
WorkflowEventSource,
@@ -16,7 +18,6 @@ from ._events import (
)
from ._runner_context import Message, RunnerContext
from ._shared_state import SharedState
from ._telemetry import workflow_tracer
T_Out = TypeVar("T_Out")
@@ -83,13 +84,19 @@ 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.
"""
global OTEL_SETTINGS
from ..observability import OTEL_SETTINGS
# Create publishing span (inherits current trace context automatically)
with workflow_tracer.create_sending_span(type(message).__name__, target_id) as span:
attributes: dict[str, str] = {OtelAttr.MESSAGE_TYPE: type(message).__name__}
if target_id:
attributes[OtelAttr.MESSAGE_DESTINATION_EXECUTOR_ID] = target_id
with create_workflow_span(OtelAttr.MESSAGE_SEND_SPAN, attributes, kind=SpanKind.PRODUCER) 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():
if OTEL_SETTINGS.ENABLED and span and span.is_recording(): # type: ignore[name-defined]
trace_context: dict[str, str] = {}
inject(trace_context) # Inject current trace context for message propagation
@@ -38,7 +38,7 @@ from .._types import (
UsageDetails,
)
from ..exceptions import ServiceInitializationError
from ..telemetry import use_telemetry
from ..observability import use_observability
from ._shared import OpenAIConfigMixin, OpenAISettings
if sys.version_info >= (3, 11):
@@ -51,7 +51,7 @@ __all__ = ["OpenAIAssistantsClient"]
@use_function_invocation
@use_telemetry
@use_observability
class OpenAIAssistantsClient(OpenAIConfigMixin, BaseChatClient):
"""OpenAI Assistants client."""
@@ -40,7 +40,7 @@ from ..exceptions import (
ServiceInvalidRequestError,
ServiceResponseException,
)
from ..telemetry import use_telemetry
from ..observability import use_observability
from ._exceptions import OpenAIContentFilterException
from ._shared import OpenAIBase, OpenAIConfigMixin, OpenAISettings, prepare_function_call_results
@@ -451,7 +451,7 @@ TOpenAIChatClient = TypeVar("TOpenAIChatClient", bound="OpenAIChatClient")
@use_function_invocation
@use_telemetry
@use_observability
class OpenAIChatClient(OpenAIConfigMixin, OpenAIBaseChatClient):
"""OpenAI Chat completion class."""
@@ -62,7 +62,7 @@ from ..exceptions import (
ServiceInvalidRequestError,
ServiceResponseException,
)
from ..telemetry import use_telemetry
from ..observability import use_observability
from ._exceptions import OpenAIContentFilterException
from ._shared import OpenAIBase, OpenAIConfigMixin, OpenAISettings, prepare_function_call_results
@@ -933,7 +933,7 @@ TOpenAIResponsesClient = TypeVar("TOpenAIResponsesClient", bound="OpenAIResponse
@use_function_invocation
@use_telemetry
@use_observability
class OpenAIResponsesClient(OpenAIConfigMixin, OpenAIBaseResponsesClient):
"""OpenAI Responses client class."""
@@ -22,9 +22,9 @@ from pydantic.types import StringConstraints
from .._logging import get_logger
from .._pydantic import AFBaseModel, AFBaseSettings
from .._telemetry import APP_INFO, USER_AGENT_KEY, prepend_agent_framework_to_user_agent
from .._types import ChatOptions, Contents, SpeechToTextOptions, TextToSpeechOptions
from ..exceptions import ServiceInitializationError
from ..telemetry import APP_INFO, USER_AGENT_KEY, prepend_agent_framework_to_user_agent
logger: logging.Logger = get_logger("agent_framework.openai")
+69
View File
@@ -0,0 +1,69 @@
# Copyright (c) Microsoft. All rights reserved.
from collections.abc import Generator
from typing import Any
from unittest.mock import patch
from opentelemetry.sdk.trace.export import SimpleSpanProcessor, SpanExporter
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
from pytest import fixture
@fixture
def enable_otel(request: Any) -> bool:
"""Fixture that returns a boolean indicating if Otel is enabled."""
return request.param if hasattr(request, "param") else True
@fixture
def enable_sensitive_data(request: Any) -> bool:
"""Fixture that returns a boolean indicating if sensitive data is enabled."""
return request.param if hasattr(request, "param") else True
@fixture(autouse=True)
def span_exporter(monkeypatch, enable_otel: bool, enable_sensitive_data: bool) -> Generator[SpanExporter]:
"""Fixture to remove environment variables for OtelSettings."""
env_vars = [
"ENABLE_OTEL",
"ENABLE_SENSITIVE_DATA",
"OTLP_ENDPOINT",
"APPLICATIONINSIGHTS_CONNECTION_STRING",
"APPLICATIONINSIGHTS_LIVE_METRICS",
]
for key in env_vars:
monkeypatch.delenv(key, raising=False) # type: ignore
monkeypatch.setenv("ENABLE_OTEL", str(enable_otel)) # type: ignore
if not enable_otel:
# we overwrite sensitive data for tests
enable_sensitive_data = False
monkeypatch.setenv("ENABLE_SENSITIVE_DATA", str(enable_sensitive_data)) # type: ignore
import importlib
from opentelemetry import trace
import agent_framework.observability as observability
# Reload the module to ensure a clean state for tests, then create a
# fresh OtelSettings instance and patch the module attribute.
importlib.reload(observability)
# recreate otel settings with values from above and no file.
otel = observability.OtelSettings(env_file_path="test.env")
otel.setup_observability()
monkeypatch.setattr(observability, "OTEL_SETTINGS", otel, raising=False) # type: ignore
exporter = InMemorySpanExporter()
with (
patch("agent_framework.observability.OTEL_SETTINGS", otel),
patch("agent_framework.observability.setup_observability"),
):
if enable_otel or enable_sensitive_data:
trace.get_tracer_provider().add_span_processor(
SimpleSpanProcessor(exporter) # type: ignore[func-returns-value]
)
yield exporter
# Clean up
exporter.clear()
@@ -27,7 +27,6 @@ from agent_framework import (
ai_function,
use_function_invocation,
)
from agent_framework.telemetry import OtelSettings, setup_telemetry
if sys.version_info >= (3, 12):
from typing import override # type: ignore
@@ -38,29 +37,6 @@ else:
logger = logging.getLogger(__name__)
@fixture
def enable_otel(request: Any) -> bool:
"""Fixture that returns a boolean indicating if Otel is enabled."""
return request.param if hasattr(request, "param") else True
@fixture
def enable_sensitive_data(request: Any) -> bool:
"""Fixture that returns a boolean indicating if sensitive data is enabled."""
return request.param if hasattr(request, "param") else False
@fixture
def otel_settings(enable_otel: bool, enable_sensitive_data: bool) -> OtelSettings:
"""Fixture to set environment variables for OtelSettings."""
from agent_framework.telemetry import OTEL_SETTINGS
setup_telemetry(enable_otel=enable_otel, enable_sensitive_data=enable_sensitive_data)
return OTEL_SETTINGS
@fixture(scope="function")
def chat_history() -> list[ChatMessage]:
return []
@@ -0,0 +1,469 @@
# Copyright (c) Microsoft. All rights reserved.
import logging
from collections.abc import MutableSequence
from typing import Any
from unittest.mock import MagicMock, Mock, patch
import pytest
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
from opentelemetry.semconv_ai import SpanAttributes
from opentelemetry.trace import StatusCode
from agent_framework import (
AGENT_FRAMEWORK_USER_AGENT,
AgentProtocol,
AgentRunResponse,
AgentThread,
BaseChatClient,
ChatMessage,
ChatOptions,
ChatResponse,
ChatResponseUpdate,
Role,
UsageDetails,
prepend_agent_framework_to_user_agent,
)
from agent_framework.exceptions import AgentInitializationError, ChatClientInitializationError
from agent_framework.observability import (
OPEN_TELEMETRY_AGENT_MARKER,
OPEN_TELEMETRY_CHAT_CLIENT_MARKER,
ROLE_EVENT_MAP,
ChatMessageListTimestampFilter,
OtelAttr,
get_function_span,
use_agent_observability,
use_observability,
)
# region Test constants
def test_role_event_map():
"""Test that ROLE_EVENT_MAP contains expected mappings."""
assert ROLE_EVENT_MAP["system"] == OtelAttr.SYSTEM_MESSAGE
assert ROLE_EVENT_MAP["user"] == OtelAttr.USER_MESSAGE
assert ROLE_EVENT_MAP["assistant"] == OtelAttr.ASSISTANT_MESSAGE
assert ROLE_EVENT_MAP["tool"] == OtelAttr.TOOL_MESSAGE
def test_enum_values():
"""Test that OtelAttr enum has expected values."""
assert OtelAttr.OPERATION == "gen_ai.operation.name"
assert SpanAttributes.LLM_SYSTEM == "gen_ai.system"
assert SpanAttributes.LLM_REQUEST_MODEL == "gen_ai.request.model"
assert OtelAttr.CHAT_COMPLETION_OPERATION == "chat"
assert OtelAttr.TOOL_EXECUTION_OPERATION == "execute_tool"
assert OtelAttr.AGENT_INVOKE_OPERATION == "invoke_agent"
# region Test ChatMessageListTimestampFilter
def test_filter_without_index_key():
"""Test filter method when record doesn't have INDEX_KEY."""
log_filter = ChatMessageListTimestampFilter()
record = logging.LogRecord(
name="test", level=logging.INFO, pathname="", lineno=0, msg="test message", args=(), exc_info=None
)
original_created = record.created
result = log_filter.filter(record)
assert result is True
assert record.created == original_created
def test_filter_with_index_key():
"""Test filter method when record has INDEX_KEY."""
log_filter = ChatMessageListTimestampFilter()
record = logging.LogRecord(
name="test", level=logging.INFO, pathname="", lineno=0, msg="test message", args=(), exc_info=None
)
original_created = record.created
# Add the index key
setattr(record, ChatMessageListTimestampFilter.INDEX_KEY, 5)
result = log_filter.filter(record)
assert result is True
# Should increment by 5 microseconds (5 * 1e-6)
assert record.created == original_created + 5 * 1e-6
def test_index_key_constant():
"""Test that INDEX_KEY constant is correctly defined."""
assert ChatMessageListTimestampFilter.INDEX_KEY == "chat_message_index"
# region Test get_function_span
def test_start_span_basic(span_exporter: InMemorySpanExporter):
"""Test starting a span with basic function info."""
# Create a mock function
mock_function = Mock()
mock_function.name = "test_function"
mock_function.description = "Test function description"
attributes = {
OtelAttr.OPERATION: OtelAttr.TOOL_EXECUTION_OPERATION,
OtelAttr.TOOL_NAME: "test_function",
OtelAttr.TOOL_DESCRIPTION: "Test function description",
OtelAttr.TOOL_TYPE: "function",
}
span_exporter.clear()
with get_function_span(attributes) as function_span:
assert function_span is not None
function_span.set_attribute("test_attr", "test_value")
spans = span_exporter.get_finished_spans()
assert len(spans) == 1
span = spans[0]
assert span.name == "execute_tool test_function"
assert span.attributes["test_attr"] == "test_value"
assert span.attributes[OtelAttr.OPERATION.value] == OtelAttr.TOOL_EXECUTION_OPERATION
assert span.attributes[OtelAttr.TOOL_NAME] == "test_function"
assert span.attributes[OtelAttr.TOOL_DESCRIPTION] == "Test function description"
def test_start_span_with_tool_call_id(span_exporter: InMemorySpanExporter):
"""Test starting a span with tool_call_id."""
tool_call_id = "test_call_123"
attributes = {
OtelAttr.OPERATION: OtelAttr.TOOL_EXECUTION_OPERATION,
OtelAttr.TOOL_NAME: "test_function",
OtelAttr.TOOL_DESCRIPTION: "Test function",
OtelAttr.TOOL_TYPE: "function",
OtelAttr.TOOL_CALL_ID: tool_call_id,
}
span_exporter.clear()
with get_function_span(attributes) as function_span:
assert function_span is not None
function_span.set_attribute("test_attr", "test_value")
spans = span_exporter.get_finished_spans()
assert len(spans) == 1
span = spans[0]
assert span.name == "execute_tool test_function"
assert span.attributes["test_attr"] == "test_value"
assert span.attributes[OtelAttr.TOOL_CALL_ID] == tool_call_id
# Verify all attributes
assert span.attributes[OtelAttr.OPERATION.value] == OtelAttr.TOOL_EXECUTION_OPERATION
assert span.attributes[OtelAttr.TOOL_NAME] == "test_function"
assert span.attributes[OtelAttr.TOOL_DESCRIPTION] == "Test function"
assert span.attributes[OtelAttr.TOOL_TYPE] == "function"
# region Test use_observability decorator
def test_decorator_with_valid_class():
"""Test that decorator works with a valid BaseChatClient-like class."""
# Create a mock class with the required methods
class MockChatClient:
async def get_response(self, messages, **kwargs):
return Mock()
async def get_streaming_response(self, messages, **kwargs):
async def gen():
yield Mock()
return gen()
# Apply the decorator
decorated_class = use_observability(MockChatClient)
assert hasattr(decorated_class, OPEN_TELEMETRY_CHAT_CLIENT_MARKER)
def test_decorator_with_missing_methods():
"""Test that decorator handles classes missing required methods gracefully."""
class MockChatClient:
OTEL_PROVIDER_NAME = "test_provider"
# Apply the decorator - should not raise an error
with pytest.raises(ChatClientInitializationError):
use_observability(MockChatClient)
def test_decorator_with_partial_methods():
"""Test decorator when only one method is present."""
class MockChatClient:
OTEL_PROVIDER_NAME = "test_provider"
async def get_response(self, messages, **kwargs):
return Mock()
with pytest.raises(ChatClientInitializationError):
use_observability(MockChatClient)
# region Test telemetry decorator with mock client
@pytest.fixture
def mock_chat_client():
"""Create a mock chat client for testing."""
class MockChatClient(BaseChatClient):
def service_url(self):
return "https://test.example.com"
async def _inner_get_response(
self, *, messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
):
return ChatResponse(
messages=[ChatMessage(role=Role.ASSISTANT, text="Test response")],
usage_details=UsageDetails(input_token_count=10, output_token_count=20),
finish_reason=None,
)
async def _inner_get_streaming_response(
self, *, messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
):
yield ChatResponseUpdate(text="Hello", role=Role.ASSISTANT)
yield ChatResponseUpdate(text=" world", role=Role.ASSISTANT)
return MockChatClient
@pytest.mark.parametrize("enable_sensitive_data", [True, False], indirect=True)
async def test_chat_client_observability(mock_chat_client, span_exporter: InMemorySpanExporter, enable_sensitive_data):
"""Test that when diagnostics are enabled, telemetry is applied."""
client = use_observability(mock_chat_client)()
messages = [ChatMessage(role=Role.USER, text="Test message")]
span_exporter.clear()
response = await client.get_response(messages=messages, ai_model_id="Test")
assert response is not None
spans = span_exporter.get_finished_spans()
assert len(spans) == 1
span = spans[0]
assert span.name == "chat Test"
assert span.attributes[OtelAttr.OPERATION.value] == OtelAttr.CHAT_COMPLETION_OPERATION
assert span.attributes[SpanAttributes.LLM_REQUEST_MODEL] == "Test"
assert span.attributes[OtelAttr.INPUT_TOKENS] == 10
assert span.attributes[OtelAttr.OUTPUT_TOKENS] == 20
if enable_sensitive_data:
assert span.attributes[OtelAttr.INPUT_MESSAGES] is not None
assert span.attributes[OtelAttr.OUTPUT_MESSAGES] is not None
@pytest.mark.parametrize("enable_sensitive_data", [True, False], indirect=True)
async def test_chat_client_streaming_observability(
mock_chat_client, span_exporter: InMemorySpanExporter, enable_sensitive_data
):
"""Test streaming telemetry through the use_observability decorator."""
client = use_observability(mock_chat_client)()
messages = [ChatMessage(role=Role.USER, text="Test")]
span_exporter.clear()
# Collect all yielded updates
updates = []
async for update in client.get_streaming_response(messages=messages, ai_model_id="Test"):
updates.append(update)
# Verify we got the expected updates, this shouldn't be dependent on otel
assert len(updates) == 2
spans = span_exporter.get_finished_spans()
assert len(spans) == 1
span = spans[0]
assert span.name == "chat Test"
assert span.attributes[OtelAttr.OPERATION.value] == OtelAttr.CHAT_COMPLETION_OPERATION
assert span.attributes[SpanAttributes.LLM_REQUEST_MODEL] == "Test"
if enable_sensitive_data:
assert span.attributes[OtelAttr.INPUT_MESSAGES] is not None
assert span.attributes[OtelAttr.OUTPUT_MESSAGES] is not None
def test_prepend_user_agent_with_none_value():
"""Test prepend user agent with None value in headers."""
headers = {"User-Agent": None}
result = prepend_agent_framework_to_user_agent(headers)
# Should handle None gracefully
assert "User-Agent" in result
assert AGENT_FRAMEWORK_USER_AGENT in str(result["User-Agent"])
# region Test use_agent_observability decorator
def test_agent_decorator_with_valid_class():
"""Test that agent decorator works with a valid ChatAgent-like class."""
# Create a mock class with the required methods
class MockChatClientAgent:
AGENT_SYSTEM_NAME = "test_agent_system"
def __init__(self):
self.id = "test_agent_id"
self.name = "test_agent"
self.display_name = "Test Agent"
self.description = "Test agent description"
async def run(self, messages=None, *, thread=None, **kwargs):
return Mock()
async def run_stream(self, messages=None, *, thread=None, **kwargs):
async def gen():
yield Mock()
return gen()
def get_new_thread(self) -> AgentThread:
return AgentThread()
# Apply the decorator
decorated_class = use_agent_observability(MockChatClientAgent)
assert hasattr(decorated_class, OPEN_TELEMETRY_AGENT_MARKER)
def test_agent_decorator_with_missing_methods():
"""Test that agent decorator handles classes missing required methods gracefully."""
class MockAgent:
AGENT_SYSTEM_NAME = "test_agent_system"
# Apply the decorator - should not raise an error
with pytest.raises(AgentInitializationError):
use_agent_observability(MockAgent)
def test_agent_decorator_with_partial_methods():
"""Test agent decorator when only one method is present."""
from agent_framework.observability import use_agent_observability
class MockAgent:
AGENT_SYSTEM_NAME = "test_agent_system"
def __init__(self):
self.id = "test_agent_id"
self.name = "test_agent"
self.display_name = "Test Agent"
async def run(self, messages=None, *, thread=None, **kwargs):
return Mock()
with pytest.raises(AgentInitializationError):
use_agent_observability(MockAgent)
# region Test agent telemetry decorator with mock agent
@pytest.fixture
def mock_chat_agent():
"""Create a mock chat client agent for testing."""
class MockChatClientAgent:
AGENT_SYSTEM_NAME = "test_agent_system"
def __init__(self):
self.id = "test_agent_id"
self.name = "test_agent"
self.display_name = "Test Agent"
self.description = "Test agent description"
async def run(self, messages=None, *, thread=None, **kwargs):
return AgentRunResponse(
messages=[ChatMessage(role=Role.ASSISTANT, text="Agent response")],
usage_details=UsageDetails(input_token_count=15, output_token_count=25),
response_id="test_response_id",
raw_representation=Mock(finish_reason=Mock(value="stop")),
)
async def run_stream(self, messages=None, *, thread=None, **kwargs):
from agent_framework import AgentRunResponseUpdate
yield AgentRunResponseUpdate(text="Hello", role=Role.ASSISTANT)
yield AgentRunResponseUpdate(text=" from agent", role=Role.ASSISTANT)
return MockChatClientAgent
@pytest.mark.parametrize("enable_sensitive_data", [True, False], indirect=True)
async def test_agent_instrumentation_enabled(
mock_chat_agent: AgentProtocol, span_exporter: InMemorySpanExporter, enable_sensitive_data
):
"""Test that when agent diagnostics are enabled, telemetry is applied."""
agent = use_agent_observability(mock_chat_agent)()
span_exporter.clear()
response = await agent.run("Test message")
assert response is not None
spans = span_exporter.get_finished_spans()
assert len(spans) == 1
span = spans[0]
assert span.name == "invoke_agent Test Agent"
assert span.attributes[OtelAttr.OPERATION.value] == OtelAttr.AGENT_INVOKE_OPERATION
assert span.attributes[OtelAttr.AGENT_ID] == "test_agent_id"
assert span.attributes[OtelAttr.AGENT_NAME] == "Test Agent"
assert span.attributes[OtelAttr.AGENT_DESCRIPTION] == "Test agent description"
assert span.attributes[SpanAttributes.LLM_REQUEST_MODEL] == "unknown"
assert span.attributes[OtelAttr.INPUT_TOKENS] == 15
assert span.attributes[OtelAttr.OUTPUT_TOKENS] == 25
if enable_sensitive_data:
assert span.attributes[OtelAttr.OUTPUT_MESSAGES] is not None
@pytest.mark.parametrize("enable_sensitive_data", [True, False], indirect=True)
async def test_agent_streaming_response_with_diagnostics_enabled_via_decorator(
mock_chat_agent: AgentProtocol, span_exporter: InMemorySpanExporter, enable_sensitive_data
):
"""Test agent streaming telemetry through the use_agent_observability decorator."""
agent = use_agent_observability(mock_chat_agent)()
span_exporter.clear()
updates = []
async for update in agent.run_stream("Test message"):
updates.append(update)
# Verify we got the expected updates
assert len(updates) == 2
spans = span_exporter.get_finished_spans()
assert len(spans) == 1
span = spans[0]
assert span.name == "invoke_agent Test Agent"
assert span.attributes[OtelAttr.OPERATION.value] == OtelAttr.AGENT_INVOKE_OPERATION
assert span.attributes[OtelAttr.AGENT_ID] == "test_agent_id"
assert span.attributes[OtelAttr.AGENT_NAME] == "Test Agent"
assert span.attributes[OtelAttr.AGENT_DESCRIPTION] == "Test agent description"
assert span.attributes[SpanAttributes.LLM_REQUEST_MODEL] == "unknown"
if enable_sensitive_data:
assert span.attributes.get(OtelAttr.OUTPUT_MESSAGES) is not None # Streaming, so no usage yet
async def test_agent_run_with_exception_handling(mock_chat_agent: AgentProtocol):
"""Test agent run with exception handling."""
async def run_with_error(self, messages=None, *, thread=None, **kwargs):
raise RuntimeError("Agent run error")
mock_chat_agent.run = run_with_error
agent = use_agent_observability(mock_chat_agent)()
from opentelemetry.trace import Span
with (
patch("agent_framework.observability._get_span") as mock_get_span,
):
mock_span = MagicMock(spec=Span)
# Ensure the patched context manager returns mock_span when entered
mock_get_span.return_value.__enter__.return_value = mock_span
# Should raise the exception and call error handler
with pytest.raises(RuntimeError, match="Agent run error"):
await agent.run("Test message")
# Verify error was recorded
# Check that both error attributes were set on the span
mock_span.set_attribute.assert_called_with(OtelAttr.ERROR_TYPE, "RuntimeError")
mock_span.record_exception.assert_called_once()
mock_span.set_status.assert_called_once_with(
status=StatusCode.ERROR, description=repr(RuntimeError("Agent run error"))
)
@@ -1,44 +1,14 @@
# Copyright (c) Microsoft. All rights reserved.
import logging
from collections.abc import MutableSequence
from typing import Any
from unittest.mock import MagicMock, Mock, patch
import pytest
from opentelemetry.semconv_ai import SpanAttributes
from opentelemetry.trace import StatusCode
from unittest.mock import patch
from agent_framework import (
AgentProtocol,
AgentRunResponse,
AgentThread,
BaseChatClient,
ChatMessage,
ChatOptions,
ChatResponse,
ChatResponseUpdate,
Role,
UsageDetails,
)
from agent_framework.exceptions import AgentInitializationError, ChatClientInitializationError
from agent_framework.telemetry import (
AGENT_FRAMEWORK_USER_AGENT,
OPEN_TELEMETRY_AGENT_MARKER,
OPEN_TELEMETRY_CHAT_CLIENT_MARKER,
ROLE_EVENT_MAP,
USER_AGENT_KEY,
USER_AGENT_TELEMETRY_DISABLED_ENV_VAR,
ChatMessageListTimestampFilter,
OtelAttr,
get_function_span,
prepend_agent_framework_to_user_agent,
use_agent_telemetry,
use_telemetry,
)
from .utils import CopyingMock
# region Test constants
@@ -59,13 +29,13 @@ def test_agent_framework_user_agent_format():
def test_app_info_when_telemetry_enabled():
"""Test that APP_INFO is set when telemetry is enabled."""
with patch("agent_framework.telemetry.IS_TELEMETRY_ENABLED", True):
with patch("agent_framework._telemetry.IS_TELEMETRY_ENABLED", True):
import importlib
import agent_framework.telemetry
import agent_framework._telemetry
importlib.reload(agent_framework.telemetry)
from agent_framework.telemetry import APP_INFO
importlib.reload(agent_framework._telemetry)
from agent_framework import APP_INFO
assert APP_INFO is not None
assert "agent-framework-version" in APP_INFO
@@ -75,7 +45,7 @@ def test_app_info_when_telemetry_enabled():
def test_app_info_when_telemetry_disabled():
"""Test that APP_INFO is None when telemetry is disabled."""
# Test the logic directly since APP_INFO is set at module import time
with patch("agent_framework.telemetry.IS_TELEMETRY_ENABLED", False):
with patch("agent_framework._telemetry.IS_TELEMETRY_ENABLED", False):
# Simulate the module's logic for APP_INFO
test_app_info = (
{
@@ -87,24 +57,6 @@ def test_app_info_when_telemetry_disabled():
assert test_app_info is None
def test_role_event_map():
"""Test that ROLE_EVENT_MAP contains expected mappings."""
assert ROLE_EVENT_MAP["system"] == OtelAttr.SYSTEM_MESSAGE
assert ROLE_EVENT_MAP["user"] == OtelAttr.USER_MESSAGE
assert ROLE_EVENT_MAP["assistant"] == OtelAttr.ASSISTANT_MESSAGE
assert ROLE_EVENT_MAP["tool"] == OtelAttr.TOOL_MESSAGE
def test_enum_values():
"""Test that OtelAttr enum has expected values."""
assert OtelAttr.OPERATION == "gen_ai.operation.name"
assert SpanAttributes.LLM_SYSTEM == "gen_ai.system"
assert SpanAttributes.LLM_REQUEST_MODEL == "gen_ai.request.model"
assert OtelAttr.CHAT_COMPLETION_OPERATION == "chat"
assert OtelAttr.TOOL_EXECUTION_OPERATION == "execute_tool"
assert OtelAttr.AGENT_INVOKE_OPERATION == "invoke_agent"
# region Test prepend_agent_framework_to_user_agent
@@ -144,415 +96,3 @@ def test_modifies_original_dict():
assert result is headers # Same object
assert "User-Agent" in headers
# region Test ChatMessageListTimestampFilter
def test_filter_without_index_key():
"""Test filter method when record doesn't have INDEX_KEY."""
log_filter = ChatMessageListTimestampFilter()
record = logging.LogRecord(
name="test", level=logging.INFO, pathname="", lineno=0, msg="test message", args=(), exc_info=None
)
original_created = record.created
result = log_filter.filter(record)
assert result is True
assert record.created == original_created
def test_filter_with_index_key():
"""Test filter method when record has INDEX_KEY."""
log_filter = ChatMessageListTimestampFilter()
record = logging.LogRecord(
name="test", level=logging.INFO, pathname="", lineno=0, msg="test message", args=(), exc_info=None
)
original_created = record.created
# Add the index key
setattr(record, ChatMessageListTimestampFilter.INDEX_KEY, 5)
result = log_filter.filter(record)
assert result is True
# Should increment by 5 microseconds (5 * 1e-6)
assert record.created == original_created + 5 * 1e-6
def test_index_key_constant():
"""Test that INDEX_KEY constant is correctly defined."""
assert ChatMessageListTimestampFilter.INDEX_KEY == "chat_message_index"
# region Test get_function_span
def test_start_span_basic():
"""Test starting a span with basic function info."""
mock_tracer = Mock()
with patch("agent_framework.telemetry.tracer", mock_tracer):
mock_span = Mock()
mock_tracer.start_as_current_span.return_value = mock_span
# Create a mock function
mock_function = Mock()
mock_function.name = "test_function"
mock_function.description = "Test function description"
attributes = {
OtelAttr.OPERATION: OtelAttr.TOOL_EXECUTION_OPERATION,
OtelAttr.TOOL_NAME: "test_function",
OtelAttr.TOOL_DESCRIPTION: "Test function description",
OtelAttr.TOOL_TYPE: "function",
}
result = get_function_span(attributes)
assert result == mock_span
mock_tracer.start_as_current_span.assert_called_once()
call_args = mock_tracer.start_as_current_span.call_args
assert call_args[1]["name"] == "execute_tool test_function"
attributes = call_args[1]["attributes"]
assert attributes[OtelAttr.OPERATION.value] == OtelAttr.TOOL_EXECUTION_OPERATION
assert attributes[OtelAttr.TOOL_NAME] == "test_function"
assert attributes[OtelAttr.TOOL_DESCRIPTION] == "Test function description"
def test_start_span_with_tool_call_id():
"""Test starting a span with tool_call_id."""
mock_tracer = Mock()
with patch("agent_framework.telemetry.tracer", mock_tracer):
mock_span = CopyingMock()
mock_tracer.start_as_current_span.return_value = mock_span
mock_function = Mock()
mock_function.name = "test_function"
mock_function.description = "Test function"
tool_call_id = "test_call_123"
attributes = {
OtelAttr.OPERATION: OtelAttr.TOOL_EXECUTION_OPERATION,
OtelAttr.TOOL_NAME: "test_function",
OtelAttr.TOOL_DESCRIPTION: "Test function",
OtelAttr.TOOL_TYPE: "function",
OtelAttr.TOOL_CALL_ID: tool_call_id,
}
_ = get_function_span(attributes)
call_args = mock_tracer.start_as_current_span.call_args
attributes = call_args[1]["attributes"]
assert attributes[OtelAttr.TOOL_CALL_ID] == "test_call_123"
# region Test use_telemetry decorator
def test_decorator_with_valid_class():
"""Test that decorator works with a valid BaseChatClient-like class."""
# Create a mock class with the required methods
class MockChatClient:
async def get_response(self, messages, **kwargs):
return Mock()
async def get_streaming_response(self, messages, **kwargs):
async def gen():
yield Mock()
return gen()
# Apply the decorator
decorated_class = use_telemetry(MockChatClient)
assert hasattr(decorated_class, OPEN_TELEMETRY_CHAT_CLIENT_MARKER)
def test_decorator_with_missing_methods():
"""Test that decorator handles classes missing required methods gracefully."""
class MockChatClient:
OTEL_PROVIDER_NAME = "test_provider"
# Apply the decorator - should not raise an error
with pytest.raises(ChatClientInitializationError):
use_telemetry(MockChatClient)
def test_decorator_with_partial_methods():
"""Test decorator when only one method is present."""
class MockChatClient:
OTEL_PROVIDER_NAME = "test_provider"
async def get_response(self, messages, **kwargs):
return Mock()
with pytest.raises(ChatClientInitializationError):
use_telemetry(MockChatClient)
# region Test telemetry decorator with mock client
@pytest.fixture
def mock_chat_client():
"""Create a mock chat client for testing."""
class MockChatClient(BaseChatClient):
def service_url(self):
return "https://test.example.com"
async def _inner_get_response(
self, *, messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
):
return ChatResponse(
messages=[ChatMessage(role=Role.ASSISTANT, text="Test response")],
usage_details=UsageDetails(input_token_count=10, output_token_count=20),
finish_reason=None,
)
async def _inner_get_streaming_response(
self, *, messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
):
yield ChatResponseUpdate(text="Hello", role=Role.ASSISTANT)
yield ChatResponseUpdate(text=" world", role=Role.ASSISTANT)
return MockChatClient
@pytest.mark.parametrize("enable_sensitive_data", [True, False], indirect=True)
async def test_instrumentation_enabled(mock_chat_client, otel_settings):
"""Test that when diagnostics are enabled, telemetry is applied."""
client = use_telemetry(mock_chat_client)()
messages = [ChatMessage(role=Role.USER, text="Test message")]
chat_options = ChatOptions()
with (
patch("agent_framework.telemetry._get_span") as mock_response_span,
patch("agent_framework.telemetry._capture_messages") as mock_log_messages,
):
response = await client.get_response(messages=messages, chat_options=chat_options)
assert response is not None
mock_response_span.assert_called_once()
# Check that log messages was called only if sensitive events are enabled
assert mock_log_messages.call_count == (2 if otel_settings.enable_sensitive_data else 0)
@pytest.mark.parametrize("enable_sensitive_data", [True, False], indirect=True)
async def test_streaming_response_with_otel(mock_chat_client, otel_settings):
"""Test streaming telemetry through the use_telemetry decorator."""
client = use_telemetry(mock_chat_client)()
messages = [ChatMessage(role=Role.USER, text="Test")]
chat_options = ChatOptions()
with (
patch("agent_framework.telemetry._get_span") as mock_response_span,
patch("agent_framework.telemetry._capture_messages") as mock_log_messages,
patch("agent_framework.telemetry._capture_response") as mock_set_output,
):
# Collect all yielded updates
updates = []
async for update in client.get_streaming_response(messages=messages, chat_options=chat_options):
updates.append(update)
# Verify we got the expected updates, this shouldn't be dependent on otel
assert len(updates) == 2
# Verify telemetry calls were made
mock_response_span.assert_called_once()
if otel_settings.enable_sensitive_data:
mock_log_messages.assert_called()
assert mock_log_messages.call_count == 2 # One for input, one for output
else:
mock_log_messages.assert_not_called()
mock_set_output.assert_called_once()
def test_prepend_user_agent_with_none_value():
"""Test prepend user agent with None value in headers."""
headers = {"User-Agent": None}
result = prepend_agent_framework_to_user_agent(headers)
# Should handle None gracefully
assert "User-Agent" in result
assert AGENT_FRAMEWORK_USER_AGENT in str(result["User-Agent"])
# region Test use_agent_telemetry decorator
def test_agent_decorator_with_valid_class():
"""Test that agent decorator works with a valid ChatAgent-like class."""
# Create a mock class with the required methods
class MockChatClientAgent:
AGENT_SYSTEM_NAME = "test_agent_system"
def __init__(self):
self.id = "test_agent_id"
self.name = "test_agent"
self.display_name = "Test Agent"
self.description = "Test agent description"
async def run(self, messages=None, *, thread=None, **kwargs):
return Mock()
async def run_stream(self, messages=None, *, thread=None, **kwargs):
async def gen():
yield Mock()
return gen()
def get_new_thread(self) -> AgentThread:
return AgentThread()
# Apply the decorator
decorated_class = use_agent_telemetry(MockChatClientAgent)
assert hasattr(decorated_class, OPEN_TELEMETRY_AGENT_MARKER)
def test_agent_decorator_with_missing_methods():
"""Test that agent decorator handles classes missing required methods gracefully."""
class MockAgent:
AGENT_SYSTEM_NAME = "test_agent_system"
# Apply the decorator - should not raise an error
with pytest.raises(AgentInitializationError):
use_agent_telemetry(MockAgent)
def test_agent_decorator_with_partial_methods():
"""Test agent decorator when only one method is present."""
from agent_framework.telemetry import use_agent_telemetry
class MockAgent:
AGENT_SYSTEM_NAME = "test_agent_system"
def __init__(self):
self.id = "test_agent_id"
self.name = "test_agent"
self.display_name = "Test Agent"
async def run(self, messages=None, *, thread=None, **kwargs):
return Mock()
with pytest.raises(AgentInitializationError):
use_agent_telemetry(MockAgent)
# region Test agent telemetry decorator with mock agent
@pytest.fixture
def mock_chat_client_agent():
"""Create a mock chat client agent for testing."""
class MockChatClientAgent:
AGENT_SYSTEM_NAME = "test_agent_system"
def __init__(self):
self.id = "test_agent_id"
self.name = "test_agent"
self.display_name = "Test Agent"
self.description = "Test agent description"
async def run(self, messages=None, *, thread=None, **kwargs):
return AgentRunResponse(
messages=[ChatMessage(role=Role.ASSISTANT, text="Agent response")],
usage_details=UsageDetails(input_token_count=15, output_token_count=25),
response_id="test_response_id",
raw_representation=Mock(finish_reason=Mock(value="stop")),
)
async def run_stream(self, messages=None, *, thread=None, **kwargs):
from agent_framework import AgentRunResponseUpdate
yield AgentRunResponseUpdate(text="Hello", role=Role.ASSISTANT)
yield AgentRunResponseUpdate(text=" from agent", role=Role.ASSISTANT)
return MockChatClientAgent
@pytest.mark.parametrize("enable_sensitive_data", [True, False], indirect=True)
async def test_agent_instrumentation_enabled(mock_chat_client_agent: AgentProtocol, otel_settings):
"""Test that when agent diagnostics are enabled, telemetry is applied."""
agent = use_agent_telemetry(mock_chat_client_agent)()
with (
patch("agent_framework.telemetry.use_span") as mock_use_span,
patch("agent_framework.telemetry.logger") as mock_logger,
):
response = await agent.run("Test message")
assert response is not None
mock_use_span.assert_called_once()
# Check that logger.info was called (telemetry logs input/output)
assert mock_logger.info.call_count == (2 if otel_settings.enable_sensitive_data else 0)
@pytest.mark.parametrize("enable_sensitive_data", [True, False], indirect=True)
async def test_agent_streaming_response_with_diagnostics_enabled_via_decorator(
mock_chat_client_agent: AgentProtocol, otel_settings
):
"""Test agent streaming telemetry through the use_agent_telemetry decorator."""
agent = use_agent_telemetry(mock_chat_client_agent)()
with (
patch("agent_framework.telemetry._get_span") as mock_get_span,
patch("agent_framework.telemetry._capture_messages") as mock_capture_messages,
patch("agent_framework.telemetry._capture_response") as mock_capture_response,
):
# Collect all yielded updates
updates = []
async for update in agent.run_stream("Test message"):
updates.append(update)
# Verify we got the expected updates
assert len(updates) == 2
# Verify telemetry calls were made
mock_get_span.assert_called_once()
mock_capture_response.assert_called_once()
if otel_settings.enable_sensitive_data:
mock_capture_messages.assert_called()
else:
mock_capture_messages.assert_not_called()
async def test_agent_run_with_exception_handling(mock_chat_client_agent: AgentProtocol):
"""Test agent run with exception handling."""
async def run_with_error(self, messages=None, *, thread=None, **kwargs):
raise RuntimeError("Agent run error")
mock_chat_client_agent.run = run_with_error
agent = use_agent_telemetry(mock_chat_client_agent)()
from opentelemetry.trace import Span
with (
patch("agent_framework.telemetry._get_span") as mock_get_span,
):
mock_span = MagicMock(spec=Span)
# Ensure the patched context manager returns mock_span when entered
mock_get_span.return_value.__enter__.return_value = mock_span
# Should raise the exception and call error handler
with pytest.raises(RuntimeError, match="Agent run error"):
await agent.run("Test message")
# Verify error was recorded
# Check that both error attributes were set on the span
mock_span.set_attribute.assert_called_with(OtelAttr.ERROR_TYPE, "RuntimeError")
mock_span.record_exception.assert_called_once()
mock_span.set_status.assert_called_once_with(
status=StatusCode.ERROR, description=repr(RuntimeError("Agent run error"))
)
+141 -149
View File
@@ -1,9 +1,10 @@
# Copyright (c) Microsoft. All rights reserved.
from typing import Any
from unittest.mock import Mock, patch
from unittest.mock import Mock
import pytest
from opentelemetry import trace
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
from pydantic import BaseModel
from agent_framework import (
@@ -15,9 +16,7 @@ from agent_framework import (
)
from agent_framework._tools import _parse_inputs
from agent_framework.exceptions import ToolException
from agent_framework.telemetry import OtelAttr
from .utils import CopyingMock
from agent_framework.observability import OtelAttr
# region AIFunction and ai_function decorator tests
@@ -85,8 +84,7 @@ async def test_ai_function_decorator_with_async():
assert (await async_test_tool(1, 2)) == 3
@pytest.mark.parametrize("enable_sensitive_data", [True], indirect=True)
async def test_ai_function_invoke_telemetry_enabled(otel_settings):
async def test_ai_function_invoke_telemetry_enabled(span_exporter: InMemorySpanExporter):
"""Test the ai_function invoke method with telemetry enabled."""
@ai_function(
@@ -97,52 +95,83 @@ async def test_ai_function_invoke_telemetry_enabled(otel_settings):
"""A function that adds two numbers for telemetry testing."""
return x + y
# Mock the tracer and span
with (
patch("agent_framework.telemetry.tracer"),
# the span creation uses a form of deepcopy, so need to mock that way
patch("agent_framework._tools.get_function_span", new_callable=CopyingMock) as mock_start_span,
):
mock_span = Mock()
mock_context_manager = Mock()
mock_context_manager.__enter__ = Mock(return_value=mock_span)
mock_context_manager.__exit__ = Mock(return_value=None)
mock_start_span.return_value = mock_context_manager
# Mock the histogram
mock_histogram = Mock()
telemetry_test_tool._invocation_duration_histogram = mock_histogram
span_exporter.clear()
# Call invoke
result = await telemetry_test_tool.invoke(x=1, y=2, tool_call_id="test_call_id")
# Mock the histogram
mock_histogram = Mock()
telemetry_test_tool._invocation_duration_histogram = mock_histogram
# Verify result
assert result == 3
# Call invoke
result = await telemetry_test_tool.invoke(x=1, y=2, tool_call_id="test_call_id")
# Verify telemetry calls
spans = span_exporter.get_finished_spans()
assert len(spans) == 1
span = spans[0]
assert OtelAttr.TOOL_EXECUTION_OPERATION.value in span.name
assert "telemetry_test_tool" in span.name
assert span.attributes[OtelAttr.TOOL_NAME] == "telemetry_test_tool"
assert span.attributes[OtelAttr.TOOL_CALL_ID] == "test_call_id"
assert span.attributes[OtelAttr.TOOL_TYPE] == "function"
assert span.attributes[OtelAttr.TOOL_DESCRIPTION] == "A test tool for telemetry"
assert span.attributes[OtelAttr.TOOL_ARGUMENTS] == '{"x": 1, "y": 2}'
assert span.attributes[OtelAttr.TOOL_RESULT] == "3"
# Verify result
assert result == 3
# Verify telemetry calls
mock_start_span.assert_called_once_with(
attributes={
OtelAttr.OPERATION: OtelAttr.TOOL_EXECUTION_OPERATION,
OtelAttr.TOOL_NAME: "telemetry_test_tool",
OtelAttr.TOOL_CALL_ID: "test_call_id",
OtelAttr.TOOL_TYPE: "function",
OtelAttr.TOOL_DESCRIPTION: "A test tool for telemetry",
OtelAttr.TOOL_ARGUMENTS: '{"x": 1, "y": 2}',
}
)
assert mock_span.set_attribute.call_count == 2
# Verify histogram was called with correct attributes
mock_histogram.record.assert_called_once()
call_args = mock_histogram.record.call_args
assert call_args[0][0] > 0 # duration should be positive
attributes = call_args[1]["attributes"]
assert attributes[OtelAttr.MEASUREMENT_FUNCTION_TAG_NAME] == "telemetry_test_tool"
assert attributes[OtelAttr.TOOL_CALL_ID] == "test_call_id"
# Verify histogram was called with correct attributes
mock_histogram.record.assert_called_once()
call_args = mock_histogram.record.call_args
assert call_args[0][0] > 0 # duration should be positive
attributes = call_args[1]["attributes"]
assert attributes[OtelAttr.MEASUREMENT_FUNCTION_TAG_NAME] == "telemetry_test_tool"
assert attributes[OtelAttr.TOOL_CALL_ID] == "test_call_id"
@pytest.mark.parametrize("enable_sensitive_data", [True], indirect=True)
async def test_ai_function_invoke_telemetry_with_pydantic_args(otel_settings):
@pytest.mark.parametrize("enable_sensitive_data", [False], indirect=True)
async def test_ai_function_invoke_telemetry_sensitive_disabled(span_exporter: InMemorySpanExporter):
"""Test the ai_function invoke method with telemetry enabled."""
@ai_function(
name="telemetry_test_tool",
description="A test tool for telemetry",
)
def telemetry_test_tool(x: int, y: int) -> int:
"""A function that adds two numbers for telemetry testing."""
return x + y
# Mock the histogram
mock_histogram = Mock()
telemetry_test_tool._invocation_duration_histogram = mock_histogram
span_exporter.clear()
# Call invoke
result = await telemetry_test_tool.invoke(x=1, y=2, tool_call_id="test_call_id")
# Verify result
assert result == 3
# Verify telemetry calls
spans = span_exporter.get_finished_spans()
assert len(spans) == 1
span = spans[0]
assert OtelAttr.TOOL_EXECUTION_OPERATION.value in span.name
assert "telemetry_test_tool" in span.name
assert span.attributes[OtelAttr.TOOL_NAME] == "telemetry_test_tool"
assert span.attributes[OtelAttr.TOOL_CALL_ID] == "test_call_id"
assert span.attributes[OtelAttr.TOOL_TYPE] == "function"
assert span.attributes[OtelAttr.TOOL_DESCRIPTION] == "A test tool for telemetry"
assert OtelAttr.TOOL_ARGUMENTS not in span.attributes
assert OtelAttr.TOOL_RESULT not in span.attributes
# Verify histogram was called with correct attributes
mock_histogram.record.assert_called_once()
call_args = mock_histogram.record.call_args
assert call_args[0][0] > 0 # duration should be positive
attributes = call_args[1]["attributes"]
assert attributes[OtelAttr.MEASUREMENT_FUNCTION_TAG_NAME] == "telemetry_test_tool"
assert attributes[OtelAttr.TOOL_CALL_ID] == "test_call_id"
async def test_ai_function_invoke_telemetry_with_pydantic_args(span_exporter: InMemorySpanExporter):
"""Test the ai_function invoke method with Pydantic model arguments."""
@ai_function(
@@ -156,42 +185,27 @@ async def test_ai_function_invoke_telemetry_with_pydantic_args(otel_settings):
# Create arguments as Pydantic model instance
args_model = pydantic_test_tool.input_model(x=5, y=10)
with (
patch("agent_framework.telemetry.tracer"),
# the span creation uses a form of deepcopy, so need to mock that way
patch("agent_framework._tools.get_function_span", new_callable=CopyingMock) as mock_start_span,
):
mock_span = Mock()
mock_context_manager = Mock()
mock_context_manager.__enter__ = Mock(return_value=mock_span)
mock_context_manager.__exit__ = Mock(return_value=None)
mock_start_span.return_value = mock_context_manager
mock_histogram = Mock()
pydantic_test_tool._invocation_duration_histogram = mock_histogram
span_exporter.clear()
# Call invoke with Pydantic model
result = await pydantic_test_tool.invoke(arguments=args_model, tool_call_id="pydantic_call")
mock_histogram = Mock()
pydantic_test_tool._invocation_duration_histogram = mock_histogram
# Call invoke with Pydantic model
result = await pydantic_test_tool.invoke(arguments=args_model, tool_call_id="pydantic_call")
# Verify result
assert result == 15
# Verify telemetry calls
mock_start_span.assert_called_once_with(
attributes={
OtelAttr.OPERATION: OtelAttr.TOOL_EXECUTION_OPERATION,
OtelAttr.TOOL_NAME: "pydantic_test_tool",
OtelAttr.TOOL_CALL_ID: "pydantic_call",
OtelAttr.TOOL_TYPE: "function",
OtelAttr.TOOL_DESCRIPTION: "A test tool with Pydantic args",
OtelAttr.TOOL_ARGUMENTS: '{"x":5,"y":10}',
}
)
assert mock_span.set_attribute.call_count == 2
# Verify result
assert result == 15
spans = span_exporter.get_finished_spans()
assert len(spans) == 1
span = spans[0]
assert OtelAttr.TOOL_EXECUTION_OPERATION.value in span.name
assert "pydantic_test_tool" in span.name
assert span.attributes[OtelAttr.TOOL_NAME] == "pydantic_test_tool"
assert span.attributes[OtelAttr.TOOL_CALL_ID] == "pydantic_call"
assert span.attributes[OtelAttr.TOOL_TYPE] == "function"
assert span.attributes[OtelAttr.TOOL_DESCRIPTION] == "A test tool with Pydantic args"
assert span.attributes[OtelAttr.TOOL_ARGUMENTS] == '{"x":5,"y":10}'
@pytest.mark.parametrize("otel_settings", [(True, True)], indirect=True)
async def test_ai_function_invoke_telemetry_with_exception(otel_settings):
async def test_ai_function_invoke_telemetry_with_exception(span_exporter: InMemorySpanExporter):
"""Test the ai_function invoke method with telemetry when an exception occurs."""
@ai_function(
@@ -202,41 +216,33 @@ async def test_ai_function_invoke_telemetry_with_exception(otel_settings):
"""A function that raises an exception for telemetry testing."""
raise ValueError("Test exception for telemetry")
with (
patch("agent_framework.telemetry.tracer"),
# the span creation uses a form of deepcopy, so need to mock that way
patch("agent_framework._tools.get_function_span", new_callable=CopyingMock) as mock_start_span,
):
mock_span = Mock()
mock_context_manager = Mock()
mock_context_manager.__enter__ = Mock(return_value=mock_span)
mock_context_manager.__exit__ = Mock(return_value=None)
mock_start_span.return_value = mock_context_manager
mock_histogram = Mock()
exception_test_tool._invocation_duration_histogram = mock_histogram
span_exporter.clear()
# Call invoke and expect exception
with pytest.raises(ValueError, match="Test exception for telemetry"):
await exception_test_tool.invoke(x=1, y=2, tool_call_id="exception_call")
spans = span_exporter.get_finished_spans()
assert len(spans) == 1
span = spans[0]
assert OtelAttr.TOOL_EXECUTION_OPERATION.value in span.name
assert "exception_test_tool" in span.name
assert span.attributes[OtelAttr.TOOL_NAME] == "exception_test_tool"
assert span.attributes[OtelAttr.TOOL_CALL_ID] == "exception_call"
assert span.attributes[OtelAttr.TOOL_TYPE] == "function"
assert span.attributes[OtelAttr.TOOL_DESCRIPTION] == "A test tool that raises an exception"
assert span.attributes[OtelAttr.TOOL_ARGUMENTS] == '{"x": 1, "y": 2}'
assert span.attributes[OtelAttr.ERROR_TYPE] == ValueError.__name__
assert span.status.status_code == trace.StatusCode.ERROR
mock_histogram = Mock()
exception_test_tool._invocation_duration_histogram = mock_histogram
# Call invoke and expect exception
with pytest.raises(ValueError, match="Test exception for telemetry"):
await exception_test_tool.invoke(x=1, y=2, tool_call_id="exception_call")
# Verify telemetry calls
mock_start_span.assert_called_once()
# Verify span exception recording
mock_span.record_exception.assert_called_once()
mock_span.set_attribute.assert_called()
mock_span.set_status.assert_called_once()
# Verify histogram was called with error attributes
mock_histogram.record.assert_called_once()
call_args = mock_histogram.record.call_args
attributes = call_args[1]["attributes"]
assert attributes[OtelAttr.ERROR_TYPE] == ValueError.__name__
# Verify histogram was called with error attributes
mock_histogram.record.assert_called_once()
call_args = mock_histogram.record.call_args
attributes = call_args[1]["attributes"]
assert attributes[OtelAttr.ERROR_TYPE] == ValueError.__name__
@pytest.mark.parametrize("enable_sensitive_data", [True], indirect=True)
async def test_ai_function_invoke_telemetry_async_function(otel_settings):
async def test_ai_function_invoke_telemetry_async_function(span_exporter: InMemorySpanExporter):
"""Test the ai_function invoke method with telemetry on async function."""
@ai_function(
@@ -247,44 +253,30 @@ async def test_ai_function_invoke_telemetry_async_function(otel_settings):
"""An async function for telemetry testing."""
return x * y
with (
patch("agent_framework.telemetry.tracer"),
# the span creation uses a form of deepcopy, so need to mock that way
patch("agent_framework._tools.get_function_span", new_callable=CopyingMock) as mock_start_span,
):
mock_span = Mock()
mock_context_manager = Mock()
mock_context_manager.__enter__ = Mock(return_value=mock_span)
mock_context_manager.__exit__ = Mock(return_value=None)
mock_start_span.return_value = mock_context_manager
mock_histogram = Mock()
async_telemetry_test._invocation_duration_histogram = mock_histogram
span_exporter.clear()
# Call invoke
result = await async_telemetry_test.invoke(x=3, y=4, tool_call_id="async_call")
mock_histogram = Mock()
async_telemetry_test._invocation_duration_histogram = mock_histogram
# Verify result
assert result == 12
spans = span_exporter.get_finished_spans()
assert len(spans) == 1
span = spans[0]
assert OtelAttr.TOOL_EXECUTION_OPERATION.value in span.name
assert "async_telemetry_test" in span.name
assert span.attributes[OtelAttr.TOOL_NAME] == "async_telemetry_test"
assert span.attributes[OtelAttr.TOOL_CALL_ID] == "async_call"
assert span.attributes[OtelAttr.TOOL_TYPE] == "function"
assert span.attributes[OtelAttr.TOOL_DESCRIPTION] == "An async test tool for telemetry"
assert span.attributes[OtelAttr.TOOL_ARGUMENTS] == '{"x": 3, "y": 4}'
# Call invoke
result = await async_telemetry_test.invoke(x=3, y=4, tool_call_id="async_call")
# Verify result
assert result == 12
# Verify telemetry calls
mock_start_span.assert_called_once_with(
attributes={
OtelAttr.OPERATION: OtelAttr.TOOL_EXECUTION_OPERATION,
OtelAttr.TOOL_NAME: "async_telemetry_test",
OtelAttr.TOOL_CALL_ID: "async_call",
OtelAttr.TOOL_TYPE: "function",
OtelAttr.TOOL_DESCRIPTION: "An async test tool for telemetry",
OtelAttr.TOOL_ARGUMENTS: '{"x": 3, "y": 4}',
}
)
assert mock_span.set_attribute.call_count == 2
# Verify histogram recording
mock_histogram.record.assert_called_once()
call_args = mock_histogram.record.call_args
attributes = call_args[1]["attributes"]
assert attributes[OtelAttr.MEASUREMENT_FUNCTION_TAG_NAME] == "async_telemetry_test"
# Verify histogram recording
mock_histogram.record.assert_called_once()
call_args = mock_histogram.record.call_args
attributes = call_args[1]["attributes"]
assert attributes[OtelAttr.MEASUREMENT_FUNCTION_TAG_NAME] == "async_telemetry_test"
async def test_ai_function_invoke_invalid_pydantic_args():
@@ -3,8 +3,6 @@ from typing import Any
from pytest import fixture
from agent_framework.telemetry import OtelSettings, setup_telemetry
# region Connector Settings fixtures
@fixture
@@ -51,26 +49,3 @@ def openai_unit_test_env(monkeypatch, exclude_list, override_env_param_dict): #
monkeypatch.setenv(key, value) # type: ignore
return env_vars
@fixture
def enable_otel(request: Any) -> bool:
"""Fixture that returns a boolean indicating if Otel is enabled."""
return request.param if hasattr(request, "param") else True
@fixture
def enable_sensitive_data(request: Any) -> bool:
"""Fixture that returns a boolean indicating if sensitive data is enabled."""
return request.param if hasattr(request, "param") else False
@fixture
def otel_settings(enable_otel: bool, enable_sensitive_data: bool) -> OtelSettings:
"""Fixture to set environment variables for OtelSettings."""
from agent_framework.telemetry import OTEL_SETTINGS
setup_telemetry(enable_otel=enable_otel, enable_sensitive_data=enable_sensitive_data)
return OTEL_SETTINGS
@@ -783,7 +783,7 @@ def test_create_streaming_response_content_with_mcp_approval_request() -> None:
@pytest.mark.parametrize("enable_otel", [False], indirect=True)
@pytest.mark.parametrize("enable_sensitive_data", [False], indirect=True)
def test_end_to_end_mcp_approval_flow(otel_settings) -> None:
def test_end_to_end_mcp_approval_flow() -> None:
"""End-to-end mocked test:
model issues an mcp_approval_request, user approves, client sends mcp_approval_response.
"""
@@ -0,0 +1 @@
# Copyright (c) Microsoft. All rights reserved.
@@ -5,9 +5,6 @@ from typing import Any
from unittest.mock import patch
import pytest
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 import (
Executor,
@@ -27,54 +24,7 @@ from agent_framework._workflow._edge import (
SwitchCaseEdgeGroupDefault,
)
from agent_framework._workflow._edge_runner import create_edge_runner
from agent_framework._workflow._telemetry import EdgeGroupDeliveryStatus, workflow_tracer
@pytest.fixture
def tracing_enabled():
"""Enable tracing for tests."""
import os
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):
"""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
from agent_framework.observability import EdgeGroupDeliveryStatus
@dataclass
@@ -1,66 +1,22 @@
# 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 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")
os.environ["AGENT_FRAMEWORK_WORKFLOW_ENABLE_OTEL"] = "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", None)
else:
os.environ["AGENT_FRAMEWORK_WORKFLOW_ENABLE_OTEL"] = 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
from agent_framework.observability import (
OtelAttr,
create_processing_span,
create_workflow_span,
)
class MockExecutor(Executor):
@@ -142,34 +98,7 @@ class FanInAggregator(Executor):
return self._processed_messages
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")
os.environ["AGENT_FRAMEWORK_WORKFLOW_ENABLE_OTEL"] = "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", None)
else:
os.environ["AGENT_FRAMEWORK_WORKFLOW_ENABLE_OTEL"] = original_value
# Reload settings again
tracer.settings = WorkflowDiagnosticSettings()
async def test_span_creation_and_attributes(tracing_enabled: Any, span_exporter: InMemorySpanExporter) -> None:
async def test_span_creation_and_attributes(span_exporter: InMemorySpanExporter) -> None:
"""Test creation and attributes of all span types (workflow, processing, sending)."""
# Create a mock workflow object
mock_workflow = cast(
@@ -186,12 +115,22 @@ async def test_span_creation_and_attributes(tracing_enabled: Any, span_exporter:
)
# 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 create_workflow_span(
OtelAttr.WORKFLOW_RUN_SPAN,
{
OtelAttr.WORKFLOW_ID: mock_workflow.id,
},
) as workflow_span:
workflow_span.add_event(OtelAttr.WORKFLOW_STARTED)
sending_attributes = {
OtelAttr.MESSAGE_TYPE: "ResponseMessage",
OtelAttr.MESSAGE_DESTINATION_EXECUTOR_ID: "target-789",
}
with (
workflow_tracer.create_processing_span("executor-456", "TestExecutor", "TestMessage") as processing_span,
workflow_tracer.create_sending_span("ResponseMessage", "target-789") as sending_span,
create_processing_span("executor-456", "TestExecutor", "TestMessage") as processing_span,
create_workflow_span(
OtelAttr.MESSAGE_SEND_SPAN, sending_attributes, kind=trace.SpanKind.PRODUCER
) as sending_span,
):
# Verify all spans are recording
assert workflow_span is not None and workflow_span.is_recording()
@@ -205,7 +144,7 @@ async def test_span_creation_and_attributes(tracing_enabled: Any, span_exporter:
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.attributes.get(OtelAttr.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
@@ -226,12 +165,14 @@ async def test_span_creation_and_attributes(tracing_enabled: Any, span_exporter:
assert sending_span.attributes.get("message.destination_executor_id") == "target-789"
async def test_trace_context_handling(tracing_enabled: Any, span_exporter: InMemorySpanExporter) -> None:
async def test_trace_context_handling(span_exporter: InMemorySpanExporter) -> None:
"""Test trace context propagation and handling in messages and executors."""
shared_state = SharedState()
ctx = InProcRunnerContext()
executor = MockExecutor("test-executor")
span_exporter.clear()
# Test trace context propagation in messages
workflow_ctx: WorkflowContext[str] = WorkflowContext(
"test-executor",
@@ -273,7 +214,8 @@ async def test_trace_context_handling(tracing_enabled: Any, span_exporter: InMem
assert processing_span.attributes.get("message.type") == "str"
async def test_trace_context_disabled_when_tracing_disabled() -> None:
@pytest.mark.parametrize("enable_otel", [False], indirect=True)
async def test_trace_context_disabled_when_tracing_disabled(enable_otel, span_exporter: InMemorySpanExporter) -> None:
"""Test that no trace context is added when tracing is disabled."""
# Tracing should be disabled by default
shared_state = SharedState()
@@ -298,7 +240,7 @@ async def test_trace_context_disabled_when_tracing_disabled() -> None:
assert message.source_span_id is None
async def test_end_to_end_workflow_tracing(tracing_enabled: Any, span_exporter: InMemorySpanExporter) -> None:
async def test_end_to_end_workflow_tracing(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")
@@ -321,7 +263,7 @@ async def test_end_to_end_workflow_tracing(tracing_enabled: Any, span_exporter:
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(OtelAttr.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(by_alias=True)
@@ -422,7 +364,7 @@ async def test_end_to_end_workflow_tracing(tracing_enabled: Any, span_exporter:
assert len(aggregator_span.links) >= 2, f"Expected at least 2 links, got {len(aggregator_span.links)}"
async def test_workflow_error_handling_in_tracing(tracing_enabled: Any, span_exporter: InMemorySpanExporter) -> None:
async def test_workflow_error_handling_in_tracing(span_exporter: InMemorySpanExporter) -> None:
"""Test that workflow errors are properly recorded in traces."""
class FailingExecutor(Executor):
@@ -457,7 +399,8 @@ async def test_workflow_error_handling_in_tracing(tracing_enabled: Any, span_exp
assert workflow_span.status.status_code.name == "ERROR"
async def test_message_trace_context_serialization() -> None:
@pytest.mark.parametrize("enable_otel", [False], indirect=True)
async def test_message_trace_context_serialization(span_exporter: InMemorySpanExporter) -> None:
"""Test that message trace context is properly serialized/deserialized."""
ctx = InProcRunnerContext()
@@ -491,7 +434,7 @@ async def test_message_trace_context_serialization() -> None:
assert restored_msg.source_span_ids == ["span123"] # Test new format
async def test_workflow_build_error_tracing(tracing_enabled: Any, span_exporter: InMemorySpanExporter) -> None:
async def test_workflow_build_error_tracing(span_exporter: InMemorySpanExporter) -> None:
"""Test that build errors are properly recorded in build spans."""
# Test validation error by not setting start executor
@@ -5,9 +5,8 @@
# see ../../../env.example for details
# Otel specific variables
APPLICATION_INSIGHTS_CONNECTION_STRING="..."
APPLICATION_INSIGHTS_LIVE_METRICS=true
APPLICATIONINSIGHTS_CONNECTION_STRING="..."
APPLICATIONINSIGHTS_LIVE_METRICS=true
OTLP_ENDPOINT="http://localhost:4317/"
ENABLE_OTEL=true
ENABLE_SENSITIVE_DATA=true
WORKFLOW_ENABLE_OTEL=true
@@ -1,7 +1,6 @@
# Copyright (c) Microsoft. All rights reserved.
# type: ignore
import asyncio
import os
from random import randint
from typing import TYPE_CHECKING, Annotated
@@ -14,18 +13,15 @@ if TYPE_CHECKING:
"""
This is the simplest sample of using the Agent Framework with telemetry.
Since it does not create a tracer or span in the script's code, we can let the Agent Framework SDK handle everything.
If the environment variables are set correctly,
the SDK will automatically initialize telemetry and collect traces and logs.
This relies on the environment setting up the telemetry, you can test this with
by navigating to this folder and running:
uv run --env-file=zero_code.env opentelemetry-instrument python 01-zero_code.py
Check the zero_code.env file for the settings used in this example and to adapt it to your environment.
"""
if "AGENT_FRAMEWORK_ENABLE_OTEL" not in os.environ:
print("Set AGENT_FRAMEWORK_ENABLE_OTEL to enable telemetry with a OTLP endpoint.")
if "AGENT_FRAMEWORK_OTLP_ENDPOINT" not in os.environ and "AGENT_FRAMEWORK_MONITOR_CONNECTION_STRING" not in os.environ:
print("Set AGENT_FRAMEWORK_OTLP_ENDPOINT or AGENT_FRAMEWORK_MONITOR_CONNECTION_STRING to enable telemetry.")
async def get_weather(
location: Annotated[str, Field(description="The location to get the weather for.")],
) -> str:
@@ -6,11 +6,10 @@ from contextlib import suppress
from random import randint
from typing import TYPE_CHECKING, Annotated, Literal
from agent_framework import __version__, ai_function
from agent_framework import ai_function
from agent_framework.observability import get_tracer, setup_observability
from agent_framework.openai import OpenAIResponsesClient
from agent_framework.telemetry import setup_telemetry
from opentelemetry import trace
from opentelemetry.trace import SpanKind
from opentelemetry.trace.span import format_trace_id
from pydantic import Field
@@ -19,8 +18,8 @@ if TYPE_CHECKING:
"""
This sample, show how you can get telemetry from a chat client and tool.
it explicitly calls the `setup_telemetry` function to set up telemetry in order to include the overall spans,
those are defined in the main and run_* functions.
it uses the `tracer` that is configured by agent framework,
which also sets up the traces with the configured environment.
"""
@@ -60,9 +59,7 @@ async def run_chat_client(client: "ChatClientProtocol", stream: bool = False) ->
"""
scenario_name = "Chat Client Stream" if stream else "Chat Client"
tracer = trace.get_tracer("agent_framework", __version__)
with tracer.start_as_current_span(name=f"Scenario: {scenario_name}", kind=SpanKind.CLIENT):
with get_tracer().start_as_current_span(name=f"Scenario: {scenario_name}", kind=trace.SpanKind.CLIENT):
print("Running scenario:", scenario_name)
message = "What's the weather in Amsterdam and in Paris?"
print(f"User: {message}")
@@ -87,9 +84,7 @@ async def run_ai_function() -> None:
The telemetry will include information about the AI function execution
and the AI service execution.
"""
tracer = trace.get_tracer("agent_framework", __version__)
with tracer.start_as_current_span("Scenario: AI Function", kind=SpanKind.CLIENT):
with get_tracer().start_as_current_span("Scenario: AI Function", kind=trace.SpanKind.CLIENT):
print("Running scenario: AI Function")
func = ai_function(get_weather)
weather = await func.invoke(location="Amsterdam")
@@ -98,11 +93,8 @@ async def run_ai_function() -> None:
async def main(scenario: Literal["chat_client", "chat_client_stream", "ai_function", "all"] = "all"):
"""Run the selected scenario(s)."""
setup_telemetry()
tracer = trace.get_tracer("My application", __version__)
with tracer.start_as_current_span("Sample Scenario's", kind=SpanKind.CLIENT) as current_span:
setup_observability()
with get_tracer().start_as_current_span("Sample Scenario's", kind=trace.SpanKind.CLIENT) as current_span:
print(f"Trace ID: {format_trace_id(current_span.get_span_context().trace_id)}")
client = OpenAIResponsesClient()
@@ -6,11 +6,10 @@ from random import randint
from typing import Annotated
from agent_framework import HostedCodeInterpreterTool
from agent_framework.telemetry import setup_telemetry
from agent_framework_foundry import FoundryChatClient
from agent_framework.foundry import FoundryChatClient
from agent_framework.observability import get_tracer, setup_observability
from azure.ai.projects.aio import AIProjectClient
from azure.identity.aio import AzureCliCredential
from opentelemetry import trace
from opentelemetry.trace import SpanKind
from opentelemetry.trace.span import format_trace_id
from pydantic import Field
@@ -19,7 +18,7 @@ from pydantic import Field
This sample, shows you can leverage the built-in telemetry in Foundry.
It uses the Foundry client to setup the telemetry, this calls
out to Foundry for a telemetry connection strings,
and then call the setup_telemetry function in the agent framework.
and then call the setup_observability function in the agent framework.
If you want to compare with the trace sent to a generic OTLP endpoint,
switch the `use_foundry_telemetry` variable to False.
"""
@@ -51,7 +50,7 @@ async def main() -> None:
In foundry you will also see specific operations happening that are called by the Foundry implementation,
such as `create_agent`.
"""
use_foundry_telemetry = True
use_foundry_obs = True
questions = [
"What's the weather in Amsterdam and in Paris?",
"Why is the sky blue?",
@@ -63,13 +62,14 @@ async def main() -> None:
AIProjectClient(endpoint=os.environ["FOUNDRY_PROJECT_ENDPOINT"], credential=credential) as project,
FoundryChatClient(client=project, setup_tracing=False) as client,
):
if use_foundry_telemetry:
await client.setup_foundry_telemetry(enable_live_metrics=True)
if use_foundry_obs:
await client.setup_foundry_observability(enable_live_metrics=True)
else:
setup_telemetry()
setup_observability()
tracer = trace.get_tracer("agent_framework")
with tracer.start_as_current_span(name="Foundry Telemetry from Agent Framework", kind=SpanKind.CLIENT) as span:
with get_tracer().start_as_current_span(
name="Foundry Telemetry from Agent Framework", kind=SpanKind.CLIENT
) as span:
for question in questions:
print(f"{BLUE}User: {question}{RESET}")
print(f"{BLUE}Assistant: {RESET}", end="")
@@ -5,9 +5,8 @@ from random import randint
from typing import Annotated
from agent_framework import ChatAgent
from agent_framework.observability import get_tracer, setup_observability
from agent_framework.openai import OpenAIChatClient
from agent_framework.telemetry import setup_telemetry
from opentelemetry import trace
from opentelemetry.trace import SpanKind
from pydantic import Field
@@ -29,12 +28,10 @@ async def get_weather(
async def main():
# Set up the telemetry
setup_telemetry()
questions = ["What's the weather in Amsterdam?", "and in Paris, and which is better?", "Why is the sky blue?"]
tracer = trace.get_tracer("agent_framework")
with tracer.start_as_current_span("Scenario: Agent Chat", kind=SpanKind.CLIENT):
setup_observability()
with get_tracer().start_as_current_span("Scenario: Agent Chat", kind=SpanKind.CLIENT):
print("Running scenario: Agent Chat")
print("Welcome to the chat, type 'exit' to quit.")
agent = ChatAgent(
@@ -6,16 +6,16 @@ from random import randint
from typing import Annotated
from agent_framework import ChatAgent
from agent_framework.observability import get_tracer
from agent_framework_foundry import FoundryChatClient
from azure.ai.projects.aio import AIProjectClient
from azure.identity.aio import AzureCliCredential
from opentelemetry import trace
from opentelemetry.trace import SpanKind
from pydantic import Field
"""
This sample shows you can can setup telemetry with a agent from Foundry.
We once again call the `setup_foundry_telemetry` method to set up telemetry in order to include the overall spans.
We once again call the `setup_foundry_observability` method to set up telemetry in order to include the overall spans.
"""
@@ -35,12 +35,11 @@ async def main():
async with (
AzureCliCredential() as credential,
AIProjectClient(endpoint=os.environ["FOUNDRY_PROJECT_ENDPOINT"], credential=credential) as project,
# this calls `setup_foundry_telemetry` through the context manager
# this calls `setup_foundry_observability` through the context manager
FoundryChatClient(client=project) as client,
):
await client.setup_foundry_telemetry(enable_live_metrics=True)
tracer = trace.get_tracer("agent_framework")
with tracer.start_as_current_span("Single Agent Chat", kind=SpanKind.CLIENT):
await client.setup_foundry_observability(enable_live_metrics=True)
with get_tracer().start_as_current_span("Single Agent Chat", kind=SpanKind.CLIENT):
print("Running Single Agent Chat")
print("Welcome to the chat, type 'exit' to quit.")
agent = ChatAgent(
@@ -10,8 +10,7 @@ from agent_framework import (
WorkflowContext,
handler,
)
from agent_framework.telemetry import setup_telemetry
from opentelemetry import trace
from agent_framework.observability import get_tracer, setup_observability
from opentelemetry.trace import SpanKind
from opentelemetry.trace.span import format_trace_id
@@ -21,6 +20,7 @@ This sample runs a simple sequential workflow with telemetry collection,
showing telemetry collection for workflow execution, executor processing,
and message publishing between executors.
"""
tracer = get_tracer("agent_framework.workflow")
# Executors for sequential workflow
@@ -65,8 +65,6 @@ async def run_sequential_workflow() -> None:
- 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:
@@ -105,10 +103,7 @@ async def run_sequential_workflow() -> None:
async def main():
"""Run the telemetry sample with a simple sequential workflow."""
setup_telemetry()
tracer = trace.get_tracer("agent_framework")
setup_observability()
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)}")
@@ -1,10 +1,11 @@
# Agent Framework Python Telemetry
# Agent Framework Python Observability
This sample folder shows how a Python application can be configured to send Agent Framework telemetry to the Application Performance Management (APM) vendors of your choice.
This sample folder shows how a Python application can be configured to send Agent Framework observability data to the Application Performance Management (APM) vendor(s) of your choice based on the Open Telemetry standard.
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) and [Aspire Dashboard](https://learn.microsoft.com/en-us/dotnet/aspire/fundamentals/dashboard/overview?tabs=bash).
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 the console.
> **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.
Or you can use the built-in tracing module of the [AI Toolkit for VS Code](https://marketplace.visualstudio.com/items?itemName=ms-windows-ai-studio.windows-ai-studio).
> 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.
@@ -12,12 +13,13 @@ 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)
2. [AI Toolkit for VS Code](https://marketplace.visualstudio.com/items?itemName=ms-windows-ai-studio.windows-ai-studio)
3. [Python Logging](https://docs.python.org/3/library/logging.html)
4. [Observability in Python](https://www.cncf.io/blog/2022/04/22/opentelemetry-and-python-a-complete-instrumentation-guide/)
## What to expect
The Agent Framework Python SDK is designed to efficiently generate comprehensive logs, traces, and metrics throughout the flow of function execution and model invocation. This allows you to effectively monitor your AI application's performance and accurately track token consumption.
The Agent Framework Python SDK is designed to efficiently generate comprehensive logs, traces, and metrics throughout the flow of function execution and model invocation. This allows you to effectively monitor your AI application's performance and accurately track token consumption. It does so based on the Semantic Conventions for GenAI defined by OpenTelemetry, and the workflows emit their own spans to provide end-to-end visibility.
## Configuration
@@ -37,17 +39,15 @@ No additional dependencies are required to enable telemetry. The necessary packa
### Environment variables
The following environment variables can be set to configure telemetry, the first two set the basic configuration:
- AGENT_FRAMEWORK_ENABLE_OTEL=true
- AGENT_FRAMEWORK_ENABLE_SENSITIVE_DATA=true
- ENABLE_OTEL=true
- ENABLE_SENSITIVE_DATA=true
Next we need to know where to send the telemetry, for that you can use either a OTLP endpoint or a connection string for Application Insights:
- AGENT_FRAMEWORK_OTLP_ENDPOINT="<url to OTLP endpoint>"
- OTLP_ENDPOINT="<url to OTLP endpoint>"
or
- AGENT_FRAMEWORK_MONITOR_CONNECTION_STRING="<connection string>"
- APPLICATIONINSIGHTS_CONNECTION_STRING="<connection string>"
Finally, you can enable live metrics streaming to Application Insights:
- AGENT_FRAMEWORK_MONITOR_LIVE_METRICS=true
> IMPORTANT - If both OTLP endpoint and connection string are set, the connection string will take precedence and there will be no trace to the OTLP endpoint.
- APPLICATIONINSIGHTS_LIVE_METRICS=true
## Samples
This folder contains different samples demonstrating how to use telemetry in various scenarios.
@@ -56,11 +56,11 @@ This folder contains different samples demonstrating how to use telemetry in var
A simple example showing how to enable telemetry in a zero-touch scenario. When the above environment variables are set, telemetry will be automatically enabled, however since you do not define any overarching tracer, you will only see the spans for the specific calls to the chat client and tools.
### [02a](./02a-generic_chat_client.py) and [02b](./02b-foundry_chat_client.py) Chat Clients:
These two samples show how to first setup the telemetry by manually importing the `setup_telemetry` function from the `agent_framework.telemetry` module and calling it. After this is done, the trace that get's created will live in the same context as the chat client calls, allowing you to see the end-to-end flow of your application. For Foundry, there is a method in the Foundry project client to get the telemetry url for your project, the `.setup_foundry_telemetry()` method in the `FoundryChatClient` class will use this url to configure telemetry and you then do not have to import and call `setup_telemetry()` manually.
Because of the way OpenTelemetry works, you can only call `setup_telemetry()` once per application run, so make sure you do that in the right place.
These two samples show how to first setup the telemetry by manually importing the `setup_observability` function from the `agent_framework.observability` module and calling it. After this is done, the trace that get's created will live in the same context as the chat client calls, allowing you to see the end-to-end flow of your application. For Foundry, there is a method in the Foundry project client to get the azure monitor connection string for your project, the `.setup_foundry_observability()` method in the `FoundryChatClient` class will use this url to configure telemetry and you then do not have to import and call `setup_observability()` manually.
If you or some other process already configure global tracer_providers or metrics_providers, the `setup_observability()` function will not override them, but instead use the existing tracer_provider, if possible. Metrics cannot be setup this way, so if you want to use metrics, you will have to call `setup_observability()` manually, before another process.
### [03a](./03a-generic_agent.py) and [03b](./03b-foundry_agent.py) Agents:
These two samples show how to setup telemetry when using the Agent Framework's agent abstraction layer. They are similar to the chat client samples, but also show how to create an agent and invoke it. The same rules apply for setting up telemetry, you can either call `setup_telemetry()` manually, or use the `setup_foundry_telemetry()` method in the `FoundryChatClient` class.
These two samples show how to setup telemetry when using the Agent Framework's agent abstraction layer. They are similar to the chat client samples, but also show how to create an agent and invoke it. The same rules apply for setting up telemetry, you can either call `setup_observability()` manually, or use the `setup_foundry_observability()` method in the `FoundryChatClient` class.
### [04 - workflow](./04-workflow.py) Workflow:
This sample shows how to setup telemetry when using the Agent Framework's workflow execution engine. It demonstrates a simple workflow scenario with telemetry.
@@ -68,32 +68,31 @@ This sample shows how to setup telemetry when using the Agent Framework's workfl
## Running the samples
1. Open a terminal and navigate to this folder: `python/samples/getting_started/telemetry/`. This is necessary for the `.env` file to be read correctly.
1. Open a terminal and navigate to this folder: `python/samples/getting_started/observability/`. This is necessary for the `.env` file to be read correctly.
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_ENABLE_OTEL=true` to enable basic telemetry and `AGENT_FRAMEWORK_ENABLE_SENSITIVE_DATA=true` to include sensitive information like prompts and responses.
> Note that `APPLICATIONINSIGHTS_CONNECTION_STRING` and `OTLP_ENDPOINT` are optional. If you don't configure them, everything will get outputted to the console.
> Set `ENABLE_OTEL=true` to enable telemetry and `ENABLE_SENSITIVE_DATA=true` to include sensitive information like prompts and responses.
> 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.
> Set `AGENT_FRAMEWORK_WORKFLOW_ENABLE_OTEL=true` to enable workflow telemetry for the workflow samples.
3. Activate your python virtual environment, and then run `python 01-zero_code.py` or others.
> This will output the Operation/Trace ID, which can be used later for filtering.
> This will also print the Operation/Trace ID, which can be used later for filtering.
## Application Insights/Azure Monitor
### Authentication
You can connect to your Application Insights instance using a connection string. You can also authenticate using Entra ID by passing a [TokenCredential](https://learn.microsoft.com/en-us/python/api/azure-core/azure.core.credentials.tokencredential?view=azure-python) to the `setup_telemetry()` function used in the samples above.
You can connect to your Application Insights instance using a connection string. You can also authenticate using Entra ID by passing a [TokenCredential](https://learn.microsoft.com/en-us/python/api/azure-core/azure.core.credentials.tokencredential?view=azure-python) to the `setup_observability()` function used in the samples above.
```python
from azure.identity import DefaultAzureCredential
setup_telemetry(credential=DefaultAzureCredential())
setup_observability(credential=DefaultAzureCredential())
```
It is recommended to use [DefaultAzureCredential](https://learn.microsoft.com/en-us/python/api/azure-identity/azure.identity.defaultazurecredential?view=azure-python) for local development and [ManagedIdentityCredential](https://learn.microsoft.com/en-us/python/api/azure-identity/azure.identity.managedidentitycredential?view=azure-python) for production environments.
### Logs and traces
Go to your Application Insights instance, click on _Transaction search_ on the left menu. Use the operation id output by the program to search for the logs and traces associated with the operation. Click on any of the search result to view the end-to-end transaction details. Read more [here](https://learn.microsoft.com/en-us/azure/azure-monitor/app/transaction-search-and-diagnostics?tabs=transaction-search).
Go to your Application Insights instance, click on _Transaction search_ on the left menu. Use the operation id printed by the program to search for the logs and traces associated with the operation. Click on any of the search result to view the end-to-end transaction details. Read more [here](https://learn.microsoft.com/en-us/azure/azure-monitor/app/transaction-search-and-diagnostics?tabs=transaction-search).
### Metrics
@@ -103,6 +102,19 @@ Running the application once will only generate one set of measurements (for eac
Please refer to here on how to analyze metrics in [Azure Monitor](https://learn.microsoft.com/en-us/azure/azure-monitor/essentials/analyze-metrics).
### Adding exporters
You can also create exporters directly and have those added to the tracer_providers, logger_providers and metrics_providers, this is useful if you want to add a different exporter on the fly, or if you want to customize the exporter. Here is an example of how to create an OTLP exporter and add it to the observability setup:
```python
from grpc import Compression
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
from agent_framework.observability import setup_observability
exporter = OTLPSpanExporter(endpoint="your-otlp-endpoint", compression=Compression.Gzip)
setup_observability(exporters=[exporter])
```
## Logs
When you are in Azure Monitor and want to have a overall view of the span, use this query in the logs section:
@@ -154,13 +166,13 @@ This will start the dashboard with:
Make sure your `.env` file includes the OTLP endpoint:
```bash
AGENT_FRAMEWORK_OTLP_ENDPOINT=http://localhost:4317
OTLP_ENDPOINT=http://localhost:4317
```
Or set it as an environment variable when running your samples:
```bash
AGENT_FRAMEWORK_ENABLE_OTEL=true AGENT_FRAMEWORK_OTLP_ENDPOINT=http://localhost:4317 python 01-zero_code.py
ENABLE_OTEL=true OTLP_ENDPOINT=http://localhost:4317 python 01-zero_code.py
```
### Viewing telemetry data
@@ -0,0 +1,6 @@
ENABLE_OTEL=true
ENABLE_SENSITIVE_DATA=true
OTEL_SERVICE_NAME=agent_framework
OTEL_TRACES_EXPORTER=otlp
OTEL_EXPORTER_OTLP_TRACES_ENDPOINT="http://localhost:4317/"
# APPLICATIONINSIGHTS_CONNECTION_STRING=""
@@ -4,7 +4,7 @@ import asyncio
import os
from typing import Any
from agent_framework import Executor, WorkflowBuilder, WorkflowContext, handler
from agent_framework import Executor, WorkflowBuilder, WorkflowContext, get_logger, handler
"""Basic tracing workflow sample.
@@ -14,7 +14,7 @@ A minimal two executor workflow demonstrates built in OpenTelemetry spans when d
The sample raises an error if tracing is not configured.
Purpose:
- Require diagnostics by checking AGENT_FRAMEWORK_WORKFLOW_ENABLE_OTEL_DIAGNOSTICS and wiring a console exporter.
- Require diagnostics by checking ENABLE_OTEL and wiring a console exporter.
- Show the span categories produced by a simple graph:
- workflow.build (events: build.started, build.validation_completed, build.completed, edge_group.process)
- workflow.run (events: workflow.started, workflow.completed or workflow.error)
@@ -26,32 +26,26 @@ Prerequisites:
- No external services required for the workflow itself.
- To print spans to the console, install the OpenTelemetry SDK: pip install opentelemetry-sdk
- Enable diagnostics:
configure your .env file with `AGENT_FRAMEWORK_WORKFLOW_ENABLE_OTEL_DIAGNOSTICS=true` or run:
export AGENT_FRAMEWORK_WORKFLOW_ENABLE_OTEL_DIAGNOSTICS=true
configure your .env file with `ENABLE_OTEL=true` or run:
export ENABLE_OTEL=true
"""
logger = get_logger()
def _ensure_tracing_configured() -> None:
"""Fail fast unless diagnostics are enabled and the SDK is present.
If the env var is set, attach a ConsoleSpanExporter so spans print to stdout.
"""
env = os.getenv("AGENT_FRAMEWORK_WORKFLOW_ENABLE_OTEL_DIAGNOSTICS", "").lower()
env = os.getenv("ENABLE_OTEL", "").lower()
if env not in {"1", "true", "yes"}:
raise RuntimeError(
"Tracing diagnostics are disabled. Set AGENT_FRAMEWORK_WORKFLOW_ENABLE_OTEL_DIAGNOSTICS=true "
"and rerun the sample."
)
try:
from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import ConsoleSpanExporter, SimpleSpanProcessor
except Exception as exc: # pragma: no cover
raise RuntimeError("OpenTelemetry SDK not found. Install it with: pip install opentelemetry-sdk") from exc
logger.info("Tracing diagnostics are disabled in the env. Setting this manually here.")
provider = TracerProvider()
provider.add_span_processor(SimpleSpanProcessor(ConsoleSpanExporter()))
trace.set_tracer_provider(provider)
from agent_framework.observability import setup_observability
from opentelemetry.sdk.trace.export import ConsoleSpanExporter
setup_observability(exporters=[ConsoleSpanExporter()])
class StartExecutor(Executor):
+3 -1
View File
@@ -1,5 +1,5 @@
version = 1
revision = 2
revision = 3
requires-python = ">=3.10"
resolution-markers = [
"python_full_version >= '3.13' and sys_platform == 'darwin'",
@@ -293,6 +293,7 @@ dependencies = [
{ name = "agent-framework-copilotstudio", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
{ name = "agent-framework-devui", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
{ name = "agent-framework-foundry", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
{ name = "agent-framework-lab-gaia", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
{ name = "agent-framework-mem0", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
{ name = "agent-framework-redis", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
]
@@ -328,6 +329,7 @@ requires-dist = [
{ name = "agent-framework-copilotstudio", editable = "packages/copilotstudio" },
{ name = "agent-framework-devui", editable = "packages/devui" },
{ name = "agent-framework-foundry", editable = "packages/foundry" },
{ name = "agent-framework-lab-gaia", editable = "packages/lab/gaia" },
{ name = "agent-framework-mem0", editable = "packages/mem0" },
{ name = "agent-framework-redis", editable = "packages/redis" },
]