test(openai-agents): cover message and tool capture modes

Author: nagkumar91

instrumentation-genai/opentelemetry-instrumentation-openai-agents/tests/test_tracer.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import json
 import sys
 from pathlib import Path
 
@@ -22,10 +23,16 @@
     OpenAIAgentsInstrumentor,
 )
 from opentelemetry.sdk.trace import TracerProvider  # noqa: E402
-from opentelemetry.sdk.trace.export import (  # noqa: E402
-    InMemorySpanExporter,
-    SimpleSpanProcessor,
-)
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor  # noqa: E402
+
+try:  # pragma: no cover - compatibility for older SDK versions
+    from opentelemetry.sdk.trace.export import (  # type: ignore[attr-defined]  # noqa: E402
+        InMemorySpanExporter,
+    )
+except ImportError:  # pragma: no cover - fallback for newer SDK layout
+    from opentelemetry.sdk.trace.export.in_memory_span_exporter import (  # noqa: E402
+        InMemorySpanExporter,
+    )
 from opentelemetry.semconv._incubating.attributes import (  # noqa: E402
     gen_ai_attributes as GenAI,
 )
@@ -34,15 +41,28 @@
 )
 from opentelemetry.trace import SpanKind  # noqa: E402
 
+GEN_AI_INPUT_MESSAGES = getattr(
+    GenAI, "GEN_AI_INPUT_MESSAGES", "gen_ai.input.messages"
+)
+GEN_AI_OUTPUT_MESSAGES = getattr(
+    GenAI, "GEN_AI_OUTPUT_MESSAGES", "gen_ai.output.messages"
+)
+GEN_AI_TOOL_CALL_ARGUMENTS = getattr(
+    GenAI, "GEN_AI_TOOL_CALL_ARGUMENTS", "gen_ai.tool.call.arguments"
+)
+GEN_AI_TOOL_CALL_RESULT = getattr(
+    GenAI, "GEN_AI_TOOL_CALL_RESULT", "gen_ai.tool.call.result"
+)
 
-def _instrument_with_provider():
+
+def _instrument_with_provider(**instrument_kwargs):
     set_trace_processors([])
     provider = TracerProvider()
     exporter = InMemorySpanExporter()
     provider.add_span_processor(SimpleSpanProcessor(exporter))
 
     instrumentor = OpenAIAgentsInstrumentor()
-    instrumentor.instrument(tracer_provider=provider)
+    instrumentor.instrument(tracer_provider=provider, **instrument_kwargs)
 
     return instrumentor, exporter
 
@@ -108,3 +128,145 @@ def test_function_span_records_tool_attributes():
     finally:
         instrumentor.uninstrument()
         exporter.clear()
+
+
+def test_generation_span_captures_messages_by_default():
+    instrumentor, exporter = _instrument_with_provider()
+
+    try:
+        with trace("workflow"):
+            with generation_span(
+                input=[{"role": "user", "content": "hi"}],
+                output=[{"role": "assistant", "content": "hello"}],
+                model="gpt-4o-mini",
+            ):
+                pass
+
+        spans = exporter.get_finished_spans()
+        client_span = next(
+            span for span in spans if span.kind is SpanKind.CLIENT
+        )
+
+        prompt = json.loads(client_span.attributes[GEN_AI_INPUT_MESSAGES])
+        completion = json.loads(client_span.attributes[GEN_AI_OUTPUT_MESSAGES])
+
+        assert prompt == [
+            {
+                "role": "user",
+                "parts": [{"type": "text", "content": "hi"}],
+            }
+        ]
+        assert completion == [
+            {
+                "role": "assistant",
+                "parts": [{"type": "text", "content": "hello"}],
+            }
+        ]
+
+        event_names = {event.name for event in client_span.events}
+        assert "gen_ai.input" in event_names
+        assert "gen_ai.output" in event_names
+
+        input_event = next(
+            event
+            for event in client_span.events
+            if event.name == "gen_ai.input"
+        )
+        output_event = next(
+            event
+            for event in client_span.events
+            if event.name == "gen_ai.output"
+        )
+
+        assert (
+            json.loads(input_event.attributes[GEN_AI_INPUT_MESSAGES]) == prompt
+        )
+        assert (
+            json.loads(output_event.attributes[GEN_AI_OUTPUT_MESSAGES])
+            == completion
+        )
+    finally:
+        instrumentor.uninstrument()
+        exporter.clear()
+
+
+def test_capture_mode_can_be_disabled():
+    instrumentor, exporter = _instrument_with_provider(
+        capture_message_content="no_content"
+    )
+
+    try:
+        with trace("workflow"):
+            with generation_span(
+                input=[{"role": "user", "content": "hi"}],
+                output=[{"role": "assistant", "content": "hello"}],
+                model="gpt-4o-mini",
+            ):
+                pass
+
+        spans = exporter.get_finished_spans()
+        client_span = next(
+            span for span in spans if span.kind is SpanKind.CLIENT
+        )
+
+        assert GEN_AI_INPUT_MESSAGES not in client_span.attributes
+        assert GEN_AI_OUTPUT_MESSAGES not in client_span.attributes
+        for event in client_span.events:
+            assert GEN_AI_INPUT_MESSAGES not in event.attributes
+            assert GEN_AI_OUTPUT_MESSAGES not in event.attributes
+    finally:
+        instrumentor.uninstrument()
+        exporter.clear()
+
+
+def test_function_span_captures_tool_payload():
+    instrumentor, exporter = _instrument_with_provider()
+
+    try:
+        with trace("workflow"):
+            with function_span(
+                name="fetch_weather",
+                input={"city": "Paris"},
+                output={"forecast": "sunny"},
+            ):
+                pass
+
+        spans = exporter.get_finished_spans()
+        tool_span = next(
+            span for span in spans if span.kind is SpanKind.INTERNAL
+        )
+
+        arguments = json.loads(
+            tool_span.attributes[GEN_AI_TOOL_CALL_ARGUMENTS]
+        )
+        result = json.loads(tool_span.attributes[GEN_AI_TOOL_CALL_RESULT])
+
+        assert arguments == {"city": "Paris"}
+        assert result == {"forecast": "sunny"}
+
+        event_names = {event.name for event in tool_span.events}
+        assert "gen_ai.tool.arguments" in event_names
+        assert "gen_ai.tool.result" in event_names
+
+        args_event = next(
+            event
+            for event in tool_span.events
+            if event.name == "gen_ai.tool.arguments"
+        )
+        result_event = next(
+            event
+            for event in tool_span.events
+            if event.name == "gen_ai.tool.result"
+        )
+
+        assert (
+            json.loads(args_event.attributes[GEN_AI_TOOL_CALL_ARGUMENTS])
+            == arguments
+        )
+        assert (
+            json.loads(result_event.attributes[GEN_AI_TOOL_CALL_RESULT])
+            == result
+        )
+    finally:
+        instrumentor.uninstrument()
+        exporter.clear()