add Bedrock instrumentor with stream parsing 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/_botocore_patch.py

@@ -0,0 +1,186 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Botocore patch for ``bedrock-agent-runtime.InvokeAgent``.
+
+Upstream ``opentelemetry-instrumentation-botocore`` has a Bedrock Runtime
+extension (``bedrock.py``) but **no Bedrock Agent Runtime extension** as of
+release 0.61b0. Therefore, for ``InvokeAgent`` calls the upstream library
+only emits a generic AWS span with ``rpc.*`` attributes and nothing
+Bedrock-specific. This patch:
+
+1. Stamps the parent botocore span with A365 / Bedrock-native attributes
+   extracted directly from the request params and the boto3 client
+   (``inputText``, ``sessionId``, ``agentId``, ``agentAliasId``, endpoint).
+2. Wraps the ``EventStream`` returned in ``response["completion"]`` with
+   :class:`BedrockAgentStreamParser` so trace events become A365 child
+   spans.
+
+Installed via :func:`wrapt.wrap_function_wrapper` on
+``botocore.client.BaseClient._make_api_call`` and idempotent.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from typing import Any
+
+from microsoft.opentelemetry.a365.core.constants import (
+    GEN_AI_CONVERSATION_ID_KEY,
+    GEN_AI_INPUT_MESSAGES_KEY,
+    GEN_AI_OPERATION_NAME_KEY,
+    GEN_AI_PROVIDER_NAME_KEY,
+    INVOKE_AGENT_OPERATION_NAME,
+)
+from opentelemetry.trace import get_current_span, get_tracer
+
+from microsoft.opentelemetry._bedrock._stream_parser import BedrockAgentStreamParser
+from microsoft.opentelemetry._bedrock._utils import (
+    AWS_BEDROCK_PROVIDER,
+    BEDROCK_AGENT_ALIAS_ID_ATTR,
+    BEDROCK_AGENT_ID_ATTR,
+    BEDROCK_SESSION_ID_ATTR,
+)
+
+_logger = logging.getLogger(__name__)
+
+_PATCHED = False
+_TRACER_NAME = "microsoft.opentelemetry._bedrock"
+
+
+def _stamp_request_attributes(span: Any, instance: Any, params: dict) -> None:
+    """Stamp A365 + Bedrock-native attributes derived from the InvokeAgent request.
+
+    Values come exclusively from ``params`` (the operation kwargs the caller
+    passed to ``invoke_agent``) and the boto3 client's endpoint metadata --
+    nothing is invented.
+    """
+    try:
+        agent_id = params.get("agentId")
+        if agent_id:
+            span.set_attribute(BEDROCK_AGENT_ID_ATTR, str(agent_id))
+
+        alias_id = params.get("agentAliasId")
+        if alias_id:
+            span.set_attribute(BEDROCK_AGENT_ALIAS_ID_ATTR, str(alias_id))
+
+        session_id = params.get("sessionId")
+        if session_id:
+            span.set_attribute(BEDROCK_SESSION_ID_ATTR, str(session_id))
+            # A365 schema: gen_ai.conversation.id is REQUIRED. For Bedrock
+            # the session_id groups related turns, matching the
+            # conversation semantic.
+            span.set_attribute(GEN_AI_CONVERSATION_ID_KEY, str(session_id))
+            span.set_attribute("microsoft.session.id", str(session_id))
+
+        input_text = params.get("inputText")
+        if input_text:
+            # A365 schema: gen_ai.input.messages is REQUIRED on InvokeAgent.
+            span.set_attribute(
+                GEN_AI_INPUT_MESSAGES_KEY,
+                json.dumps([{"role": "user", "content": str(input_text)}]),
+            )
+
+        # server.address / server.port are REQUIRED on InvokeAgent.
+        endpoint_url = getattr(getattr(instance, "meta", None), "endpoint_url", None)
+        if endpoint_url:
+            # boto3 endpoint_url looks like "https://bedrock-agent-runtime.us-east-1.amazonaws.com"
+            from urllib.parse import urlparse
+
+            parsed = urlparse(str(endpoint_url))
+            if parsed.hostname:
+                span.set_attribute("server.address", parsed.hostname)
+            # Default HTTPS port if scheme is https and no explicit port.
+            port = parsed.port or (443 if parsed.scheme == "https" else 80 if parsed.scheme == "http" else None)
+            if port is not None:
+                span.set_attribute("server.port", port)
+
+        # Ensure operation-name + provider are tagged even before the
+        # span processor runs (defensive: the processor runs on_start
+        # before this hook fires, so this is a no-op in normal flow).
+        span.set_attribute(GEN_AI_OPERATION_NAME_KEY, INVOKE_AGENT_OPERATION_NAME)
+        span.set_attribute(GEN_AI_PROVIDER_NAME_KEY, AWS_BEDROCK_PROVIDER)
+    except Exception as exc:  # pragma: no cover - defensive
+        _logger.debug("Failed to stamp InvokeAgent request attributes: %s", exc)
+
+
+def install_botocore_patch() -> None:
+    """Install the InvokeAgent EventStream wrapper. Safe to call multiple times."""
+    global _PATCHED
+    if _PATCHED:
+        return
+    try:
+        import wrapt  # type: ignore[import-not-found]
+        import botocore.client  # type: ignore[import-not-found]
+    except ImportError:
+        _logger.debug("botocore or wrapt not installed; Bedrock stream parsing disabled.")
+        return
+
+    def _wrap_invoke_agent(wrapped: Any, instance: Any, args: tuple, kwargs: dict) -> Any:
+        service_id = ""
+        try:
+            service_id = instance.meta.service_model.service_id  # type: ignore[attr-defined]
+        except Exception:
+            service_id = ""
+
+        # botocore signature: _make_api_call(operation_name, api_params)
+        operation = args[0] if args else kwargs.get("operation_name", "")
+        if str(service_id).replace(" ", "").lower() != "bedrockagentruntime" or operation != "InvokeAgent":
+            return wrapped(*args, **kwargs)
+
+        params: dict = {}
+        if len(args) > 1 and isinstance(args[1], dict):
+            params = args[1]
+        elif "api_params" in kwargs and isinstance(kwargs["api_params"], dict):
+            params = kwargs["api_params"]
+
+        # The upstream botocore extension has already opened a span for this
+        # API call by the time this wrapper runs; grab it now so we can
+        # stamp request-side attributes before invoking the call.
+        parent_span = get_current_span()
+        if parent_span and parent_span.get_span_context().is_valid:
+            _stamp_request_attributes(parent_span, instance, params)
+
+        response = wrapped(*args, **kwargs)
+        if not isinstance(response, dict):
+            return response
+
+        stream = response.get("completion")
+        if stream is None:
+            return response
+
+        if not parent_span or not parent_span.get_span_context().is_valid:
+            return response
+
+        tracer = get_tracer(_TRACER_NAME)
+        parser = BedrockAgentStreamParser(parent_span=parent_span, tracer=tracer)
+        # Replace the EventStream with a generator wrapping it. boto3
+        # callers iterate ``response["completion"]`` lazily, so this
+        # preserves streaming semantics.
+        response["completion"] = parser.wrap(stream)
+        return response
+
+    wrapt.wrap_function_wrapper(
+        botocore.client.BaseClient,
+        "_make_api_call",
+        _wrap_invoke_agent,
+    )
+    _PATCHED = True
+
+
+def uninstall_botocore_patch() -> None:
+    """Best-effort removal of the patch (uses wrapt's unwrap)."""
+    global _PATCHED
+    if not _PATCHED:
+        return
+    try:
+        import botocore.client  # type: ignore[import-not-found]
+
+        original = getattr(botocore.client.BaseClient._make_api_call, "__wrapped__", None)
+        if original is not None:
+            botocore.client.BaseClient._make_api_call = original  # type: ignore[assignment]
+    except Exception:
+        pass
+    finally:
+        _PATCHED = False

src/microsoft/opentelemetry/_bedrock/_span_enricher.py

@@ -0,0 +1,60 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Span enricher for AWS Bedrock spans (final attribute mapping at export)."""
+
+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 (
+    AWS_BEDROCK_PROVIDER,
+    BEDROCK_AGENT_ALIAS_ID_ATTR,
+    BEDROCK_KNOWLEDGE_BASE_ID_ATTR,
+    GEN_AI_SYSTEM_ATTR,
+)
+
+
+def enrich_bedrock_span(span: ReadableSpan) -> ReadableSpan:
+    """Map Bedrock-native botocore attributes onto A365 GenAI schema keys."""
+    attrs = span.attributes or {}
+    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)
+    if provider != AWS_BEDROCK_PROVIDER and system != AWS_BEDROCK_PROVIDER:
+        return span
+
+    extra: dict[str, object] = {}
+    if GEN_AI_PROVIDER_NAME_KEY not in attrs and system:
+        extra[GEN_AI_PROVIDER_NAME_KEY] = str(system)
+
+    # 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)
+
+    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,70 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Rename Bedrock botocore spans to A365 operation names on start."""
+
+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 (
+    AWS_BEDROCK_PROVIDER,
+    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,
+)
+
+
+class BedrockSpanProcessor(SpanProcessor):
+    """Rewrite Bedrock botocore spans to A365 operation names."""
+
+    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)
+                span.update_name(
+                    f"{INVOKE_AGENT_OPERATION_NAME} {agent_id}" if agent_id else 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:
+            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)
+            span.update_name(f"{CHAT_OPERATION_NAME} {model}" if model else 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