From a19d2416a58aaae1b88dd0c035bc59791318c0ca Mon Sep 17 00:00:00 2001 From: Jianbiao Lu Date: Mon, 18 May 2026 14:22:37 -0700 Subject: [PATCH] Add Bedrock instrumentation support --- pyproject.toml | 2 + .../opentelemetry/_bedrock/__init__.py | 6 ++ .../opentelemetry/_bedrock/_span_enricher.py | 74 +++++++++++++++ .../opentelemetry/_bedrock/_span_processor.py | 90 +++++++++++++++++++ .../_bedrock/_trace_instrumentor.py | 65 ++++++++++++++ .../opentelemetry/_bedrock/_utils.py | 39 ++++++++ src/microsoft/opentelemetry/_constants.py | 2 + 7 files changed, 278 insertions(+) create mode 100644 src/microsoft/opentelemetry/_bedrock/__init__.py create mode 100644 src/microsoft/opentelemetry/_bedrock/_span_enricher.py create mode 100644 src/microsoft/opentelemetry/_bedrock/_span_processor.py create mode 100644 src/microsoft/opentelemetry/_bedrock/_trace_instrumentor.py create mode 100644 src/microsoft/opentelemetry/_bedrock/_utils.py diff --git a/pyproject.toml b/pyproject.toml index c95f63cc..f02e1a96 100644 --- a/pyproject.toml +++ b/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"} diff --git a/src/microsoft/opentelemetry/_bedrock/__init__.py b/src/microsoft/opentelemetry/_bedrock/__init__.py new file mode 100644 index 00000000..43db35a4 --- /dev/null +++ b/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"] diff --git a/src/microsoft/opentelemetry/_bedrock/_span_enricher.py b/src/microsoft/opentelemetry/_bedrock/_span_enricher.py new file mode 100644 index 00000000..3e373777 --- /dev/null +++ b/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 diff --git a/src/microsoft/opentelemetry/_bedrock/_span_processor.py b/src/microsoft/opentelemetry/_bedrock/_span_processor.py new file mode 100644 index 00000000..49d1b9fe --- /dev/null +++ b/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 diff --git a/src/microsoft/opentelemetry/_bedrock/_trace_instrumentor.py b/src/microsoft/opentelemetry/_bedrock/_trace_instrumentor.py new file mode 100644 index 00000000..94053c5b --- /dev/null +++ b/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() diff --git a/src/microsoft/opentelemetry/_bedrock/_utils.py b/src/microsoft/opentelemetry/_bedrock/_utils.py new file mode 100644 index 00000000..0939abba --- /dev/null +++ b/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() diff --git a/src/microsoft/opentelemetry/_constants.py b/src/microsoft/opentelemetry/_constants.py index 99c71fd9..4e1c7db5 100644 --- a/src/microsoft/opentelemetry/_constants.py +++ b/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