feat: add mark_agent_root() to mark a trace's agent root span

.github/workflows/ci.yml

@@ -109,7 +109,8 @@ jobs:
 
     - name: Install package
       run: |
-        python -m pip install --upgrade pip
+        # The runner's preinstalled setuptools is audited too; keep it patched.
+        python -m pip install --upgrade pip setuptools
         pip install -e ".[aws]"
 
     - name: Run pip-audit

README.md

@@ -7,6 +7,7 @@ OTel-native tracing and scoring for LLM applications. Instrument your AI workflo
 - **One-line setup** — `register()` configures the full OTel pipeline (TracerProvider, exporter, auto-instrumentation)
 - **`observe()`** — single decorator / context manager that creates OTel spans with [GenAI semantic convention](https://opentelemetry.io/docs/specs/semconv/gen-ai/) attributes
 - **`enrich()`** — add model, token usage, and other GenAI attributes to the active span from anywhere in your code
+- **`mark_agent_root()`** — mark the agent's span as the trace's agent root in OpenSearch Agent Traces
 - **Auto-instrumentation** — automatically discovers and activates installed instrumentor packages (OpenAI, Anthropic, Bedrock, LangChain, etc.)
 - **Scoring** — `score()` emits evaluation metrics as OTel spans at span, trace, or session level
 - **Benchmarks** — `evaluate()` runs your agent against a dataset with scorers; `Benchmark` uploads results from any eval framework (RAGAS, DeepEval, pytest)
@@ -53,17 +54,20 @@ from opensearch_genai_observability_sdk_py import register, observe, Op, enrich,
 # 1. Initialize tracing (one line)
 register(endpoint="http://localhost:21890/opentelemetry/v1/traces")
 
+
 # 2. Trace your functions
 @observe(name="web_search", op=Op.EXECUTE_TOOL)
 def search(query: str) -> list[dict]:
     return [{"title": f"Result for: {query}"}]
 
+
 @observe(name="research_agent", op=Op.INVOKE_AGENT)
 def research(query: str) -> str:
     results = search(query)
     enrich(model="gpt-4.1", provider="openai", input_tokens=150, output_tokens=50)
     return f"Summary of: {results}"
 
+
 # 3. Use context managers for inline blocks
 @observe(name="qa_pipeline", op=Op.INVOKE_AGENT)
 def run(question: str) -> str:
@@ -72,6 +76,7 @@ def run(question: str) -> str:
         enrich(safe=True)
     return answer
 
+
 result = run("What is OpenSearch?")
 
 # 4. Submit scores (after workflow completes)
@@ -136,7 +141,7 @@ Configures the OTel tracing pipeline. Call once at startup.
 register(
     endpoint="http://my-collector:4318/v1/traces",  # or use env vars
     service_name="my-app",
-    batch=True,            # BatchSpanProcessor (True) or Simple (False)
+    batch=True,  # BatchSpanProcessor (True) or Simple (False)
     auto_instrument=True,  # discover installed instrumentor packages
 )
 ```
@@ -192,10 +197,10 @@ def plan(query: str) -> str:
     enrich(model="gpt-4.1")
     return call_llm(query)
 
+
 # Without parentheses (uses function name, no op)
 @observe
-def my_function():
-    ...
+def my_function(): ...
 ```
 
 **As a context manager:**
@@ -279,6 +284,19 @@ def chat(prompt: str) -> str:
 | `session_id` | `gen_ai.session.id` | Session/conversation ID |
 | `**extra` | As provided | Any additional key-value attributes |
 
+### `mark_agent_root()`
+
+Mark the current span as its trace's agent root. OpenSearch Agent Traces lists a trace by the span marked with `derived.gen_ai.is_root`. Data Prepper marks the outermost GenAI span of each trace for you (the root when it is a GenAI span, or the agent span below an HTTP handler), so you only need this to choose the span yourself. A span marked at the source wins.
+
+```python
+@observe(name="planner", op=Op.INVOKE_AGENT)
+def plan(query: str) -> str:
+    mark_agent_root()
+    ...
+```
+
+Without the SDK, set the span attribute `derived.gen_ai.is_root` to `true` on the agent's span.
+
 ### `score()`
 
 Submits evaluation scores as OTel spans. Use any evaluation framework you prefer (autoevals, RAGAS, custom) and submit the results through `score()`.
@@ -332,13 +350,16 @@ Run a task against a dataset, score outputs, and record results as OTel spans. A
 ```python
 from opensearch_genai_observability_sdk_py import evaluate, EvalScore, observe, Op
 
+
 @observe(op=Op.INVOKE_AGENT)
 def my_agent(question: str) -> str:
     return call_llm(question)
 
+
 def accuracy(input, output, expected) -> EvalScore:
     return EvalScore(name="accuracy", value=1.0 if expected in output else 0.0)
 
+
 result = evaluate(
     name="rag-agent",
     task=my_agent,
@@ -431,6 +452,7 @@ roots = retriever.list_root_spans(services=["my-agent"], max_results=10)
 
 # Filter by time
 from datetime import datetime
+
 roots = retriever.list_root_spans(services=["my-agent"], since=datetime(2026, 3, 16))
 
 # Check which traces already have evaluation spans

src/opensearch_genai_observability_sdk_py/__init__.py

@@ -21,7 +21,7 @@
     TestCaseResult,
     evaluate,
 )
-from opensearch_genai_observability_sdk_py.enrich import enrich
+from opensearch_genai_observability_sdk_py.enrich import enrich, mark_agent_root
 from opensearch_genai_observability_sdk_py.exporters import AWSSigV4OTLPExporter
 from opensearch_genai_observability_sdk_py.observe import Op, observe
 from opensearch_genai_observability_sdk_py.register import register
@@ -35,6 +35,7 @@
     "observe",
     "Op",
     "enrich",
+    "mark_agent_root",
     # Scoring
     "score",
     # Benchmarks

src/opensearch_genai_observability_sdk_py/enrich.py

@@ -160,3 +160,31 @@ def enrich(
     for key, value in extra.items():
         if value is not None:
             span.set_attribute(key, value)
+
+
+AGENT_ROOT_ATTRIBUTE = "derived.gen_ai.is_root"
+"""Marks the span that OpenSearch Agent Traces lists as a trace's agent root."""
+
+
+def mark_agent_root() -> None:
+    """Mark the current span as the agent root of its trace.
+
+    OpenSearch Agent Traces lists a trace by the span marked with
+    ``derived.gen_ai.is_root``. Data Prepper marks the outermost GenAI span of
+    each trace for you: the root when it is a GenAI span, or the agent span
+    below it (for example when an HTTP handler runs the agent). Call this to
+    choose the span yourself, for example the agent's entry span in a service
+    that is called by another agent. A span marked at the source wins, and
+    Data Prepper then marks nothing else in that part of the trace.
+
+    Example::
+
+        @observe(name="planner", op=Op.INVOKE_AGENT)
+        def plan(query: str) -> str:
+            mark_agent_root()
+            ...
+    """
+    span = trace.get_current_span()
+    if not span.is_recording():
+        return
+    span.set_attribute(AGENT_ROOT_ATTRIBUTE, True)

src/opensearch_genai_observability_sdk_py/exporters.py

@@ -144,13 +144,15 @@ def __init__(
                 "AWS_DEFAULT_REGION environment variable, or ~/.aws/config."
             )
 
-        # Pass the signing session to OTLPSpanExporter.  The parent stores it
-        # as self._session and routes all HTTP calls through it.
-        kwargs["session"] = _SigV4AuthSession(
+        # Pass the signing session to OTLPSpanExporter, which routes all HTTP
+        # calls through it. Keep our own reference: where the parent stores the
+        # session is an implementation detail (1.45 moved it into a transport).
+        self.signing_session = _SigV4AuthSession(
             credentials=credentials,
             service=service,
             region=resolved_region,
         )
+        kwargs["session"] = self.signing_session
         super().__init__(*args, **kwargs)
 
         logger.info(

tests/test_enrich.py

@@ -11,7 +11,11 @@
 
 import json
 
-from opensearch_genai_observability_sdk_py.enrich import enrich
+from opensearch_genai_observability_sdk_py.enrich import (
+    AGENT_ROOT_ATTRIBUTE,
+    enrich,
+    mark_agent_root,
+)
 from opensearch_genai_observability_sdk_py.observe import Op, observe
 
 # ---------------------------------------------------------------------------
@@ -181,3 +185,28 @@ def test_enrich_without_span_does_not_crash(self, exporter):
         enrich(model="gpt-4.1", input_tokens=100)
         spans = exporter.get_finished_spans()
         assert len(spans) == 0
+
+
+# ---------------------------------------------------------------------------
+# mark_agent_root()
+# ---------------------------------------------------------------------------
+
+
+class TestMarkAgentRoot:
+    def test_marks_the_current_span(self, exporter):
+        @observe(name="planner", op=Op.INVOKE_AGENT)
+        def fn():
+            mark_agent_root()
+            with observe(name="chat", op=Op.CHAT):
+                pass
+            return "ok"
+
+        fn()
+        spans = {s.name: s for s in exporter.get_finished_spans()}
+        assert AGENT_ROOT_ATTRIBUTE == "derived.gen_ai.is_root"
+        assert spans["invoke_agent planner"].attributes[AGENT_ROOT_ATTRIBUTE] is True
+        assert AGENT_ROOT_ATTRIBUTE not in spans["chat chat"].attributes
+
+    def test_without_span_does_not_crash(self, exporter):
+        mark_agent_root()
+        assert exporter.get_finished_spans() == ()

tests/test_exporters.py

@@ -211,23 +211,42 @@ def test_explicit_region_overrides_botocore(self):
         """region= kwarg takes precedence over whatever botocore detects."""
         with self._mock_botocore(region="ap-southeast-1"):
             exporter = AWSSigV4OTLPExporter(endpoint=ENDPOINT, region="eu-west-1")
-        assert exporter._session._region == "eu-west-1"
+        assert exporter.signing_session._region == "eu-west-1"
 
     def test_session_is_sigv4_auth_session(self):
         """The exporter's internal session must be our signing session."""
         with self._mock_botocore():
             exporter = AWSSigV4OTLPExporter(endpoint=ENDPOINT)
-        assert isinstance(exporter._session, _SigV4AuthSession)
+        assert isinstance(exporter.signing_session, _SigV4AuthSession)
 
     def test_session_carries_correct_service(self):
         with self._mock_botocore():
             exporter = AWSSigV4OTLPExporter(endpoint=ENDPOINT, service="es")
-        assert exporter._session._service == "es"
+        assert exporter.signing_session._service == "es"
 
     def test_default_service_is_osis(self):
         with self._mock_botocore():
             exporter = AWSSigV4OTLPExporter(endpoint=ENDPOINT)
-        assert exporter._session._service == "osis"
+        assert exporter.signing_session._service == "osis"
+
+    def test_exports_through_the_signing_session(self):
+        """Every export goes through the SigV4 session, whatever the parent does with it."""
+        from opentelemetry.sdk.trace import TracerProvider
+        from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+
+        with self._mock_botocore():
+            exporter = AWSSigV4OTLPExporter(endpoint=ENDPOINT)
+        response = MagicMock(status_code=200, ok=True)
+        with patch.object(
+            _SigV4AuthSession, "request", autospec=True, return_value=response
+        ) as request:
+            provider = TracerProvider()
+            provider.add_span_processor(SimpleSpanProcessor(exporter))
+            provider.get_tracer("test").start_span("span").end()
+            provider.shutdown()
+
+        assert request.call_count >= 1
+        assert request.call_args.args[0] is exporter.signing_session
 
 
 # ---------------------------------------------------------------------------