Add Bedrock instrumentation support

pyproject.toml

@@ -37,6 +37,7 @@ dependencies = [
     "opentelemetry-instrumentation-urllib==0.61b0",
     "opentelemetry-instrumentation-urllib3==0.61b0",
     "opentelemetry-instrumentation-logging==0.61b0",
+    "opentelemetry-instrumentation-botocore==0.61b0",
     "opentelemetry-instrumentation-openai-agents-v2==0.1.0",
     "opentelemetry-instrumentation-openai-v2==2.3b0",
     "opentelemetry-resource-detector-azure<1.0.0,>=0.1.5",
@@ -83,6 +84,7 @@ docs = [
 langchain = "microsoft.opentelemetry._genai._langchain._tracer_instrumentor:LangChainInstrumentor"
 semantic_kernel = "microsoft.opentelemetry._semantic_kernel._trace_instrumentor:SemanticKernelInstrumentor"
 agent_framework = "microsoft.opentelemetry._agent_framework._trace_instrumentor:AgentFrameworkInstrumentor"
+bedrock = "microsoft.opentelemetry._bedrock._trace_instrumentor:BedrockInstrumentor"
 
 [tool.setuptools]
 package-dir = {"" = "src"}

src/microsoft/opentelemetry/_bedrock/__init__.py

@@ -0,0 +1,6 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from microsoft.opentelemetry._bedrock._trace_instrumentor import BedrockInstrumentor
+
+__all__ = ["BedrockInstrumentor"]

src/microsoft/opentelemetry/_bedrock/_span_enricher.py

@@ -0,0 +1,74 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Span enricher for AWS Bedrock spans.
+
+Maps Bedrock-specific botocore attributes onto the A365 / OTel GenAI schema
+keys so the Agent 365 ingestion pipeline can classify and display them
+without bespoke knowledge of AWS naming.
+"""
+
+from __future__ import annotations
+
+from microsoft.opentelemetry.a365.core.constants import (
+    EXECUTE_TOOL_OPERATION_NAME,
+    GEN_AI_AGENT_VERSION_KEY,
+    GEN_AI_OPERATION_NAME_KEY,
+    GEN_AI_PROVIDER_NAME_KEY,
+    GEN_AI_TOOL_NAME_KEY,
+    GEN_AI_TOOL_TYPE_KEY,
+)
+from microsoft.opentelemetry.a365.core.exporters.enriched_span import EnrichedReadableSpan
+from opentelemetry.sdk.trace import ReadableSpan
+
+from microsoft.opentelemetry._bedrock._utils import (
+    BEDROCK_AGENT_ALIAS_ID_ATTR,
+    BEDROCK_KNOWLEDGE_BASE_ID_ATTR,
+    GEN_AI_SYSTEM_ATTR,
+)
+
+
+def enrich_bedrock_span(span: ReadableSpan) -> ReadableSpan:
+    """Map Bedrock botocore attributes onto A365 GenAI schema keys."""
+    attrs = span.attributes or {}
+
+    # Only touch spans the BedrockSpanProcessor has tagged.
+    operation = attrs.get(GEN_AI_OPERATION_NAME_KEY)
+    provider = attrs.get(GEN_AI_PROVIDER_NAME_KEY)
+    if not operation:
+        return span
+    system = attrs.get(GEN_AI_SYSTEM_ATTR)
+    is_bedrock = provider == "aws.bedrock" or system == "aws.bedrock"
+    if not is_bedrock:
+        return span
+
+    extra: dict[str, object] = {}
+
+    # Mirror gen_ai.system -> gen_ai.provider.name if upstream only set the
+    # legacy key.
+    if GEN_AI_PROVIDER_NAME_KEY not in attrs and system:
+        extra[GEN_AI_PROVIDER_NAME_KEY] = str(system)
+
+    # Map Bedrock agent identifiers onto the A365 agent identity keys (best
+    # effort -- customers will typically also stamp tenant-scoped identities
+    # via BaggageBuilder).
+    # Map Bedrock alias_id onto gen_ai.agent.version (optional field).
+    # NOTE: Bedrock's agentId is NOT the A365 agent GUID. Per the A365
+    # schema, ``gen_ai.agent.id`` must be the customer-registered A365 GUID;
+    # we therefore do NOT map ``aws.bedrock.agent.id`` to it and leave that
+    # attribute to be populated by the application via ``BaggageBuilder``.
+    alias_id = attrs.get(BEDROCK_AGENT_ALIAS_ID_ATTR)
+    if alias_id and GEN_AI_AGENT_VERSION_KEY not in attrs:
+        extra[GEN_AI_AGENT_VERSION_KEY] = str(alias_id)
+
+    # Knowledge-base retrieve calls become datastore tool invocations.
+    if operation == EXECUTE_TOOL_OPERATION_NAME:
+        kb_id = attrs.get(BEDROCK_KNOWLEDGE_BASE_ID_ATTR)
+        if kb_id and GEN_AI_TOOL_NAME_KEY not in attrs:
+            extra[GEN_AI_TOOL_NAME_KEY] = f"bedrock-kb/{kb_id}"
+        if GEN_AI_TOOL_TYPE_KEY not in attrs:
+            extra[GEN_AI_TOOL_TYPE_KEY] = "datastore"
+
+    if extra:
+        return EnrichedReadableSpan(span, extra)
+    return span

src/microsoft/opentelemetry/_bedrock/_span_processor.py

@@ -0,0 +1,90 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""SpanProcessor that renames AWS Bedrock spans to Agent 365 operation names.
+
+The upstream ``opentelemetry-instrumentation-botocore`` package already
+creates spans for ``bedrock-runtime`` and ``bedrock-agent-runtime`` API
+calls and populates the standard ``gen_ai.*`` attributes for Converse /
+InvokeModel calls. This processor:
+
+1. Renames the span on start so the A365 ingestion classifier recognises
+   the operation (``invoke_agent`` / ``chat`` / ``execute_tool``).
+2. Sets ``gen_ai.operation.name`` and ``gen_ai.provider.name`` for
+   operations that the upstream extension does not yet cover
+   (``InvokeAgent`` and the Bedrock knowledge-base ``Retrieve`` family).
+"""
+
+from __future__ import annotations
+
+from microsoft.opentelemetry.a365.core.constants import (
+    CHAT_OPERATION_NAME,
+    EXECUTE_TOOL_OPERATION_NAME,
+    GEN_AI_OPERATION_NAME_KEY,
+    GEN_AI_PROVIDER_NAME_KEY,
+    GEN_AI_REQUEST_MODEL_KEY,
+    INVOKE_AGENT_OPERATION_NAME,
+)
+from opentelemetry import context as context_api
+from opentelemetry.sdk.trace import ReadableSpan, Span
+from opentelemetry.sdk.trace.export import SpanProcessor
+
+from microsoft.opentelemetry._bedrock._utils import (
+    BEDROCK_AGENT_ID_ATTR,
+    BEDROCK_AGENT_INVOKE_OPS,
+    BEDROCK_AGENT_RETRIEVE_OPS,
+    BEDROCK_AGENT_RUNTIME_SERVICE,
+    BEDROCK_KNOWLEDGE_BASE_ID_ATTR,
+    BEDROCK_RUNTIME_INFERENCE_OPS,
+    BEDROCK_RUNTIME_SERVICE,
+    normalize_service,
+)
+
+_AWS_BEDROCK_PROVIDER = "aws.bedrock"
+
+
+class BedrockSpanProcessor(SpanProcessor):
+    """Rewrite Bedrock botocore spans to A365 operation names on start."""
+
+    def on_start(self, span: Span, parent_context: context_api.Context | None = None) -> None:
+        attrs = span.attributes or {}
+        rpc_method = attrs.get("rpc.method")
+        if not rpc_method:
+            return
+
+        service = normalize_service(attrs.get("rpc.service"))
+        if service == BEDROCK_AGENT_RUNTIME_SERVICE:
+            if rpc_method in BEDROCK_AGENT_INVOKE_OPS:
+                span.set_attribute(GEN_AI_OPERATION_NAME_KEY, INVOKE_AGENT_OPERATION_NAME)
+                span.set_attribute(GEN_AI_PROVIDER_NAME_KEY, _AWS_BEDROCK_PROVIDER)
+                agent_id = attrs.get(BEDROCK_AGENT_ID_ATTR)
+                if agent_id:
+                    span.update_name(f"{INVOKE_AGENT_OPERATION_NAME} {agent_id}")
+                else:
+                    span.update_name(INVOKE_AGENT_OPERATION_NAME)
+            elif rpc_method in BEDROCK_AGENT_RETRIEVE_OPS:
+                span.set_attribute(GEN_AI_OPERATION_NAME_KEY, EXECUTE_TOOL_OPERATION_NAME)
+                span.set_attribute(GEN_AI_PROVIDER_NAME_KEY, _AWS_BEDROCK_PROVIDER)
+                kb_id = attrs.get(BEDROCK_KNOWLEDGE_BASE_ID_ATTR)
+                tool_label = f"bedrock-kb/{kb_id}" if kb_id else "bedrock-kb"
+                span.update_name(f"{EXECUTE_TOOL_OPERATION_NAME} {tool_label}")
+        elif service == BEDROCK_RUNTIME_SERVICE and rpc_method in BEDROCK_RUNTIME_INFERENCE_OPS:
+            # Upstream botocore already sets GEN_AI_OPERATION_NAME for these
+            # ops, but we mirror the value to be defensive and add the
+            # provider name (which upstream does not set).
+            span.set_attribute(GEN_AI_OPERATION_NAME_KEY, CHAT_OPERATION_NAME)
+            span.set_attribute(GEN_AI_PROVIDER_NAME_KEY, _AWS_BEDROCK_PROVIDER)
+            model = attrs.get(GEN_AI_REQUEST_MODEL_KEY)
+            if model:
+                span.update_name(f"{CHAT_OPERATION_NAME} {model}")
+            else:
+                span.update_name(CHAT_OPERATION_NAME)
+
+    def on_end(self, span: ReadableSpan) -> None:
+        pass
+
+    def shutdown(self) -> None:
+        pass
+
+    def force_flush(self, timeout_millis: int = 30000) -> bool:
+        return True

src/microsoft/opentelemetry/_bedrock/_trace_instrumentor.py

@@ -0,0 +1,65 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Auto-instrumentor that wires Bedrock spans into the A365 pipeline.
+
+This instrumentor does not patch botocore itself -- it relies on the
+upstream ``opentelemetry-instrumentation-botocore`` package for that.
+Instead, on activation it:
+
+* Adds a :class:`BedrockSpanProcessor` to the active TracerProvider so
+  Bedrock spans get the right A365 operation names on start.
+* Registers an enricher that maps Bedrock attributes onto the A365 GenAI
+  schema keys before export.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Collection
+from typing import Any
+
+from microsoft.opentelemetry.a365.core.exporters.enriching_span_processor import (
+    register_span_enricher,
+    unregister_span_enricher,
+)
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor  # type: ignore[attr-defined]
+from opentelemetry.trace import get_tracer_provider
+
+from microsoft.opentelemetry._bedrock._span_enricher import enrich_bedrock_span
+from microsoft.opentelemetry._bedrock._span_processor import BedrockSpanProcessor
+
+_logger = logging.getLogger(__name__)
+_instruments = ("botocore >= 1.0.0",)
+
+
+class BedrockInstrumentor(BaseInstrumentor):
+    """Instruments AWS Bedrock (via botocore) for Agent 365 observability."""
+
+    _processor: BedrockSpanProcessor | None = None
+    _owns_enricher: bool = False
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return _instruments
+
+    def _instrument(self, **kwargs: Any) -> None:
+        provider = kwargs.get("tracer_provider") or get_tracer_provider()
+
+        self._processor = BedrockSpanProcessor()
+        provider.add_span_processor(self._processor)  # type: ignore[union-attr, attr-defined]
+
+        try:
+            register_span_enricher(enrich_bedrock_span)
+            self._owns_enricher = True
+        except RuntimeError:
+            _logger.debug(
+                "A span enricher is already registered. Skipping Bedrock enricher registration."
+            )
+
+    def _uninstrument(self, **kwargs: Any) -> None:
+        if self._owns_enricher:
+            unregister_span_enricher()
+            self._owns_enricher = False
+
+        if self._processor is not None:
+            self._processor.shutdown()

src/microsoft/opentelemetry/_bedrock/_utils.py

@@ -0,0 +1,39 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Utility helpers for the AWS Bedrock observability extension."""
+
+from __future__ import annotations
+
+# AWS botocore service IDs (case-insensitive, with whitespace stripped). The
+# upstream ``opentelemetry-instrumentation-botocore`` package sets
+# ``rpc.service`` to the human-readable service id (e.g. ``"Bedrock Runtime"``,
+# ``"Bedrock Agent Runtime"``); we normalize for comparison.
+BEDROCK_RUNTIME_SERVICE = "bedrockruntime"
+BEDROCK_AGENT_RUNTIME_SERVICE = "bedrockagentruntime"
+
+BEDROCK_RUNTIME_INFERENCE_OPS = frozenset(
+    {
+        "InvokeModel",
+        "InvokeModelWithResponseStream",
+        "Converse",
+        "ConverseStream",
+    }
+)
+
+BEDROCK_AGENT_INVOKE_OPS = frozenset({"InvokeAgent"})
+
+BEDROCK_AGENT_RETRIEVE_OPS = frozenset({"Retrieve", "RetrieveAndGenerate"})
+
+# Botocore Bedrock-specific attribute keys (set by the upstream extension).
+BEDROCK_AGENT_ID_ATTR = "aws.bedrock.agent.id"
+BEDROCK_AGENT_ALIAS_ID_ATTR = "aws.bedrock.agent.alias.id"
+BEDROCK_KNOWLEDGE_BASE_ID_ATTR = "aws.bedrock.knowledge_base.id"
+GEN_AI_SYSTEM_ATTR = "gen_ai.system"
+
+
+def normalize_service(service: str | None) -> str:
+    """Lowercase and strip whitespace from an AWS rpc.service string."""
+    if not service:
+        return ""
+    return service.replace(" ", "").lower()

src/microsoft/opentelemetry/_constants.py

@@ -43,6 +43,8 @@
     "openai_agents",
     "semantic_kernel",
     "agent_framework",
+    "botocore",
+    "bedrock",
 )
 
 # Libraries disabled by default when A365 is enabled (agent workloads