Move suppress_invoke_agent_input to core SDK (exporter and span processor)

Co-authored-by: juliomenendez <9697+juliomenendez@users.noreply.github.com>

Author: Copilot

libraries/microsoft-agents-a365-observability-core/microsoft_agents_a365/observability/core/config.py

@@ -53,6 +53,7 @@ def configure(
         token_resolver: Callable[[str, str], str | None] | None = None,
         cluster_category: str = "prod",
         exporter_options: Optional[Agent365ExporterOptions] = None,
+        suppress_invoke_agent_input: bool = False,
         **kwargs: Any,
     ) -> bool:
         """
@@ -67,6 +68,7 @@ def configure(
             Use exporter_options instead.
         :param exporter_options: Agent365ExporterOptions instance for configuring the exporter.
             If provided, exporter_options takes precedence. If exporter_options is None, the token_resolver and cluster_category parameters are used as fallback/legacy support to construct a default Agent365ExporterOptions instance.
+        :param suppress_invoke_agent_input: If True, suppress input messages for spans that are children of InvokeAgent spans.
         :return: True if configuration succeeded, False otherwise.
         """
         try:
@@ -78,6 +80,7 @@ def configure(
                     token_resolver,
                     cluster_category,
                     exporter_options,
+                    suppress_invoke_agent_input,
                     **kwargs,
                 )
         except Exception as e:
@@ -92,6 +95,7 @@ def _configure_internal(
         token_resolver: Callable[[str, str], str | None] | None = None,
         cluster_category: str = "prod",
         exporter_options: Optional[Agent365ExporterOptions] = None,
+        suppress_invoke_agent_input: bool = False,
         **kwargs: Any,
     ) -> bool:
         """Internal configuration method - not thread-safe, must be called with lock."""
@@ -151,6 +155,7 @@ def _configure_internal(
                 token_resolver=exporter_options.token_resolver,
                 cluster_category=exporter_options.cluster_category,
                 use_s2s_endpoint=exporter_options.use_s2s_endpoint,
+                suppress_invoke_agent_input=suppress_invoke_agent_input,
             )
         else:
             exporter = ConsoleSpanExporter()
@@ -162,7 +167,7 @@ def _configure_internal(
 
         # Create BatchSpanProcessor with optimized settings
         batch_processor = BatchSpanProcessor(exporter, **batch_processor_kwargs)
-        agent_processor = SpanProcessor()
+        agent_processor = SpanProcessor(suppress_invoke_agent_input=suppress_invoke_agent_input)
 
         tracer_provider.add_span_processor(batch_processor)
         tracer_provider.add_span_processor(agent_processor)

libraries/microsoft-agents-a365-observability-core/microsoft_agents_a365/observability/core/exporters/agent365_exporter.py

@@ -17,6 +17,7 @@
 from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
 from opentelemetry.trace import StatusCode
 
+from ..constants import GEN_AI_INPUT_MESSAGES_KEY, GEN_AI_OPERATION_NAME_KEY, INVOKE_AGENT_OPERATION_NAME
 from .utils import (
     get_validated_domain_override,
     hex_span_id,
@@ -52,6 +53,7 @@ def __init__(
         token_resolver: Callable[[str, str], str | None],
         cluster_category: str = "prod",
         use_s2s_endpoint: bool = False,
+        suppress_invoke_agent_input: bool = False,
     ):
         if token_resolver is None:
             raise ValueError("token_resolver must be provided.")
@@ -61,6 +63,7 @@ def __init__(
         self._token_resolver = token_resolver
         self._cluster_category = cluster_category
         self._use_s2s_endpoint = use_s2s_endpoint
+        self._suppress_invoke_agent_input = suppress_invoke_agent_input
         # Read domain override once at initialization
         self._domain_override = get_validated_domain_override()
 
@@ -222,13 +225,22 @@ def _post_with_retries(self, url: str, body: str, headers: dict[str, str]) -> bo
     # ------------- Payload mapping ------------------
 
     def _build_export_request(self, spans: Sequence[ReadableSpan]) -> dict[str, Any]:
+        # Build a map of span IDs to their operation names for parent lookups
+        span_operation_map = {}
+        if self._suppress_invoke_agent_input:
+            for sp in spans:
+                attrs = sp.attributes or {}
+                operation_name = attrs.get(GEN_AI_OPERATION_NAME_KEY)
+                if operation_name:
+                    span_operation_map[sp.context.span_id] = operation_name
+        
         # Group by instrumentation scope (name, version)
         scope_map: dict[tuple[str, str | None], list[dict[str, Any]]] = {}
 
         for sp in spans:
             scope = sp.instrumentation_scope
             scope_key = (scope.name, scope.version)
-            scope_map.setdefault(scope_key, []).append(self._map_span(sp))
+            scope_map.setdefault(scope_key, []).append(self._map_span(sp, span_operation_map))
 
         scope_spans: list[dict[str, Any]] = []
         for (name, version), mapped_spans in scope_map.items():
@@ -257,7 +269,7 @@ def _build_export_request(self, spans: Sequence[ReadableSpan]) -> dict[str, Any]
             ]
         }
 
-    def _map_span(self, sp: ReadableSpan) -> dict[str, Any]:
+    def _map_span(self, sp: ReadableSpan, span_operation_map: dict[int, str] = None) -> dict[str, Any]:
         ctx = sp.context
 
         parent_span_id = None
@@ -266,6 +278,16 @@ def _map_span(self, sp: ReadableSpan) -> dict[str, Any]:
 
         # attributes
         attrs = dict(sp.attributes or {})
+        
+        # Suppress input messages if configured and parent is an InvokeAgent span
+        if self._suppress_invoke_agent_input and span_operation_map:
+            # Check if parent span is an InvokeAgent span
+            if sp.parent is not None and sp.parent.span_id != 0:
+                parent_operation = span_operation_map.get(sp.parent.span_id)
+                if parent_operation == INVOKE_AGENT_OPERATION_NAME:
+                    # Remove input messages attribute
+                    attrs.pop(GEN_AI_INPUT_MESSAGES_KEY, None)
+        
         # events
         events = []
         for ev in sp.events:

libraries/microsoft-agents-a365-observability-core/microsoft_agents_a365/observability/core/trace_processor/span_processor.py

@@ -24,8 +24,15 @@
 class SpanProcessor(BaseSpanProcessor):
     """Span processor that propagates every baggage key/value to span attributes."""
 
-    def __init__(self):
+    def __init__(self, suppress_invoke_agent_input: bool = False):
+        """Initialize the span processor.
+        
+        Args:
+            suppress_invoke_agent_input: If True, suppress input messages for spans 
+                that are children of InvokeAgent spans.
+        """
         super().__init__()
+        self._suppress_invoke_agent_input = suppress_invoke_agent_input
 
     def on_start(self, span, parent_context=None):
         ctx = parent_context or context.get_current()
@@ -81,4 +88,9 @@ def on_start(self, span, parent_context=None):
         return super().on_start(span, parent_context)
 
     def on_end(self, span):
+        """Called when a span ends.
+        
+        Note: Input suppression for InvokeAgent scopes is handled by the exporter,
+        not in this processor, because span attributes cannot be removed after they're set.
+        """
         super().on_end(span)

libraries/microsoft-agents-a365-observability-extensions-openai/microsoft_agents_a365/observability/extensions/openai/trace_instrumentor.py

@@ -47,7 +47,6 @@ def _instrument(self, **kwargs: Any) -> None:
         """Instruments the OpenAI Agents SDK with Microsoft Agent 365 tracing."""
         tracer_name = kwargs["tracer_name"] if kwargs.get("tracer_name") else None
         tracer_version = kwargs["tracer_version"] if kwargs.get("tracer_version") else None
-        suppress_invoke_agent_input = kwargs.get("suppress_invoke_agent_input", False)
 
         # Get the configured Microsoft Agent 365 Tracer
         try:
@@ -65,9 +64,7 @@ def _instrument(self, **kwargs: Any) -> None:
 
         agent365_tracer = cast(Tracer, tracer)
 
-        set_trace_processors(
-            [OpenAIAgentsTraceProcessor(agent365_tracer, suppress_invoke_agent_input)]
-        )
+        set_trace_processors([OpenAIAgentsTraceProcessor(agent365_tracer)])
 
     def _uninstrument(self, **kwargs: Any) -> None:
         pass

libraries/microsoft-agents-a365-observability-extensions-openai/microsoft_agents_a365/observability/extensions/openai/trace_processor.py

@@ -69,9 +69,8 @@
 class OpenAIAgentsTraceProcessor(TracingProcessor):
     _MAX_HANDOFFS_IN_FLIGHT = 1000
 
-    def __init__(self, tracer: Tracer, suppress_invoke_agent_input: bool = False) -> None:
+    def __init__(self, tracer: Tracer) -> None:
         self._tracer = tracer
-        self._suppress_invoke_agent_input = suppress_invoke_agent_input
         self._root_spans: dict[str, OtelSpan] = {}
         self._otel_spans: dict[str, OtelSpan] = {}
         self._tokens: dict[str, object] = {}
@@ -90,29 +89,6 @@ def _stamp_custom_parent(self, otel_span: OtelSpan, trace_id: str) -> None:
         pid_hex = "0x" + ot_trace.format_span_id(sc.span_id)
         otel_span.set_attribute(CUSTOM_PARENT_SPAN_ID_KEY, pid_hex)
 
-    def _should_suppress_input(self, span: Span[Any], otel_span: OtelSpan) -> bool:
-        """Check if input messages should be suppressed for the given span.
-        
-        Args:
-            span: The agents SDK span being processed.
-            otel_span: The corresponding OpenTelemetry span.
-            
-        Returns:
-            True if suppression is enabled and the span has an InvokeAgent parent.
-        """
-        if not self._suppress_invoke_agent_input:
-            return False
-        
-        # Check if this span has a parent that is an InvokeAgent span
-        if span.parent_id:
-            parent_otel_span = self._otel_spans.get(span.parent_id)
-            if parent_otel_span and hasattr(parent_otel_span, 'attributes'):
-                operation_name = parent_otel_span.attributes.get(GEN_AI_OPERATION_NAME_KEY)
-                if operation_name == INVOKE_AGENT_OPERATION_NAME:
-                    return True
-        
-        return False
-
     def on_trace_start(self, trace: Trace) -> None:
         """Called when a trace is started.
 
@@ -178,8 +154,7 @@ def on_span_end(self, span: Span[Any]) -> None:
                 otel_span.set_attribute(GEN_AI_OUTPUT_MESSAGES_KEY, response.model_dump_json())
                 for k, v in get_attributes_from_response(response):
                     otel_span.set_attribute(k, v)
-            # Only record input messages if not suppressing or not in InvokeAgent scope
-            if not self._should_suppress_input(span, otel_span) and hasattr(data, "input") and (input := data.input):
+            if hasattr(data, "input") and (input := data.input):
                 if isinstance(input, str):
                     otel_span.set_attribute(GEN_AI_INPUT_MESSAGES_KEY, input)
                 elif isinstance(input, list):
@@ -189,12 +164,7 @@ def on_span_end(self, span: Span[Any]) -> None:
                 elif TYPE_CHECKING:
                     assert_never(input)
         elif isinstance(data, GenerationSpanData):
-            # Collect all attributes once and filter if suppression is enabled
-            should_suppress = self._should_suppress_input(span, otel_span)
             for k, v in get_attributes_from_generation_span_data(data):
-                # Skip input messages if suppression is enabled and in InvokeAgent scope
-                if should_suppress and k == GEN_AI_INPUT_MESSAGES_KEY:
-                    continue
                 otel_span.set_attribute(k, v)
             self._stamp_custom_parent(otel_span, span.trace_id)
             otel_span.update_name(

tests/observability/extensions/openai/test_prompt_suppression.py

@@ -2,65 +2,58 @@
 
 import unittest
 
-from microsoft_agents_a365.observability.core import configure, get_tracer
-from microsoft_agents_a365.observability.extensions.openai.trace_processor import (
-    OpenAIAgentsTraceProcessor,
-)
+from microsoft_agents_a365.observability.core.exporters.agent365_exporter import _Agent365Exporter
+from microsoft_agents_a365.observability.core.trace_processor.span_processor import SpanProcessor
 
 
 class TestPromptSuppressionConfiguration(unittest.TestCase):
-    """Unit tests for prompt suppression configuration in OpenAIAgentsTraceProcessor."""
-
-    @classmethod
-    def setUpClass(cls):
-        """Set up test environment once for all tests."""
-        configure(
-            service_name="test-service-prompt-suppression",
-            service_namespace="test-namespace-prompt-suppression",
-        )
+    """Unit tests for prompt suppression configuration in the core SDK."""
 
-    def test_default_suppression_is_false(self):
-        """Test that the default value for suppress_invoke_agent_input is False."""
-        tracer = get_tracer()
-        processor = OpenAIAgentsTraceProcessor(tracer)
+    def test_span_processor_default_suppression_is_false(self):
+        """Test that the default value for suppress_invoke_agent_input is False in SpanProcessor."""
+        processor = SpanProcessor()
 
         self.assertFalse(
             processor._suppress_invoke_agent_input,
             "Default value for suppress_invoke_agent_input should be False",
         )
 
-    def test_can_enable_suppression(self):
-        """Test that suppression can be enabled via constructor."""
-        tracer = get_tracer()
-        processor = OpenAIAgentsTraceProcessor(tracer, suppress_invoke_agent_input=True)
+    def test_span_processor_can_enable_suppression(self):
+        """Test that suppression can be enabled via SpanProcessor constructor."""
+        processor = SpanProcessor(suppress_invoke_agent_input=True)
 
         self.assertTrue(
             processor._suppress_invoke_agent_input,
             "suppress_invoke_agent_input should be True when explicitly set",
         )
 
-    def test_can_disable_suppression(self):
-        """Test that suppression can be explicitly disabled via constructor."""
-        tracer = get_tracer()
-        processor = OpenAIAgentsTraceProcessor(tracer, suppress_invoke_agent_input=False)
+    def test_span_processor_can_disable_suppression(self):
+        """Test that suppression can be explicitly disabled via SpanProcessor constructor."""
+        processor = SpanProcessor(suppress_invoke_agent_input=False)
 
         self.assertFalse(
             processor._suppress_invoke_agent_input,
             "suppress_invoke_agent_input should be False when explicitly set",
         )
 
-    def test_has_should_suppress_input_method(self):
-        """Test that the processor has the helper method for suppression logic."""
-        tracer = get_tracer()
-        processor = OpenAIAgentsTraceProcessor(tracer, suppress_invoke_agent_input=True)
+    def test_exporter_default_suppression_is_false(self):
+        """Test that the default value for suppress_invoke_agent_input is False in exporter."""
+        exporter = _Agent365Exporter(token_resolver=lambda x, y: "test")
 
-        self.assertTrue(
-            hasattr(processor, "_should_suppress_input"),
-            "Processor should have _should_suppress_input method",
+        self.assertFalse(
+            exporter._suppress_invoke_agent_input,
+            "Default value for suppress_invoke_agent_input should be False",
         )
+
+    def test_exporter_can_enable_suppression(self):
+        """Test that suppression can be enabled via exporter constructor."""
+        exporter = _Agent365Exporter(
+            token_resolver=lambda x, y: "test", suppress_invoke_agent_input=True
+        )
+
         self.assertTrue(
-            callable(processor._should_suppress_input),
-            "_should_suppress_input should be callable",
+            exporter._suppress_invoke_agent_input,
+            "suppress_invoke_agent_input should be True when explicitly set",
         )