mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
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:
committed by
GitHub
Unverified
parent
f93f16a9ad
commit
2576e7a091
+6
-6
@@ -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
|
||||
@@ -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
|
||||
|
||||
|
||||
+547
-214
File diff suppressed because it is too large
Load Diff
@@ -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")
|
||||
|
||||
|
||||
@@ -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"))
|
||||
)
|
||||
|
||||
@@ -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
|
||||
|
||||
+33
-90
@@ -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
|
||||
+2
-3
@@ -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
|
||||
+6
-10
@@ -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:
|
||||
+8
-16
@@ -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()
|
||||
+10
-10
@@ -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="")
|
||||
+3
-6
@@ -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(
|
||||
+5
-6
@@ -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(
|
||||
+3
-8
@@ -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)}")
|
||||
|
||||
+36
-24
@@ -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):
|
||||
|
||||
Generated
+3
-1
@@ -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" },
|
||||
]
|
||||
|
||||
Reference in New Issue
Block a user