CCM-33842: add Bedrock GenAI span enrichment

src/harness_sdk/instrumentation/botocore/__init__.py

@@ -1,11 +1,170 @@
 '''Hypertrace wrapper around OTel botocore Instrumentor''' # pylint: disable=R0801
-import logging
+from typing import Any, Callable, Optional
+
 from opentelemetry.instrumentation.botocore import BotocoreInstrumentor # pylint:disable=E0611,E0401
+
 from harness_sdk.instrumentation import BaseInstrumentorWrapper
 
 
 from harness_sdk.custom_logger import get_custom_logger
 logger = get_custom_logger(__name__)
 
+
+_BEDROCK_RUNTIME_SERVICE_NAMES = frozenset({"bedrock-runtime", "Bedrock Runtime"})
+_BEDROCK_CONVERSE_OPERATIONS = frozenset({"Converse", "ConverseStream"})
+_BEDROCK_MODEL_ID_HEADER = "x-amzn-bedrock-model-id"
+
+
+def _is_bedrock_converse(service_name: str, operation_name: str) -> bool:
+    return (
+        service_name in _BEDROCK_RUNTIME_SERVICE_NAMES
+        and operation_name in _BEDROCK_CONVERSE_OPERATIONS
+    )
+
+
+def _get_value(source: Any, key: str, default: Any = None) -> Any:
+    if source is None:
+        return default
+    if isinstance(source, dict):
+        return source.get(key, default)
+    return getattr(source, key, default)
+
+
+def _set_if_present(span: Any, key: str, value: Any) -> None:
+    if value is None:
+        return
+    span.set_attribute(key, value)
+
+
+def _bedrock_request_model(api_params: dict[str, Any]) -> Optional[str]:
+    model_id = api_params.get("modelId")
+    return str(model_id) if model_id else None
+
+
+def _is_inference_profile_identifier(model_id: Optional[str]) -> bool:
+    return bool(
+        model_id
+        and (
+            ":inference-profile/" in model_id
+            or ":application-inference-profile/" in model_id
+            or model_id.startswith("inference-profile/")
+            or model_id.startswith("application-inference-profile/")
+        )
+    )
+
+
+def _bedrock_inference_config(api_params: dict[str, Any]) -> dict[str, Any]:
+    config = api_params.get("inferenceConfig") or {}
+    if not isinstance(config, dict):
+        return {}
+    return config
+
+
+def _set_bedrock_request_attributes(
+    span: Any,
+    operation_name: str,
+    api_params: dict[str, Any],
+) -> None:
+    if span is None or not span.is_recording():
+        return
+
+    inference_config = _bedrock_inference_config(api_params)
+    request_model = _bedrock_request_model(api_params)
+
+    _set_if_present(span, "gen_ai.request.model", request_model)
+    span.set_attribute("gen_ai.operation.name", "chat")
+    span.set_attribute("gen_ai.provider.name", "aws.bedrock")
+    span.set_attribute("gen_ai.framework", "boto3")
+    if _is_inference_profile_identifier(request_model):
+        _set_if_present(span, "aws.bedrock.inference_profile_arn", request_model)
+    span.set_attribute(
+        "gen_ai.request.streaming",
+        operation_name == "ConverseStream",
+    )
+    _set_if_present(
+        span,
+        "gen_ai.request.max_tokens",
+        inference_config.get("maxTokens"),
+    )
+    _set_if_present(
+        span,
+        "gen_ai.request.temperature",
+        inference_config.get("temperature"),
+    )
+    _set_if_present(span, "gen_ai.request.top_p", inference_config.get("topP"))
+
+
+def _bedrock_execution_model(response: Any) -> Optional[str]:
+    response_metadata = _get_value(response, "ResponseMetadata") or {}
+    headers = _get_value(response_metadata, "HTTPHeaders") or {}
+    model_id = _get_value(headers, _BEDROCK_MODEL_ID_HEADER)
+    return str(model_id) if model_id else None
+
+
+def _set_bedrock_response_attributes(
+    span: Any,
+    response: Any,
+    request_model: Optional[str],
+) -> None:
+    if span is None or not span.is_recording():
+        return
+
+    execution_model = _bedrock_execution_model(response)
+    response_model = execution_model or request_model
+    _set_if_present(span, "gen_ai.response.model", response_model)
+    _set_if_present(
+        span,
+        "aws.bedrock.execution_model_id",
+        execution_model,
+    )
+    _set_if_present(
+        span,
+        "gen_ai.response.finish_reasons",
+        _get_value(response, "stopReason"),
+    )
+
+    usage = _get_value(response, "usage") or {}
+    _set_if_present(
+        span,
+        "gen_ai.usage.input_tokens",
+        _get_value(usage, "inputTokens"),
+    )
+    _set_if_present(
+        span,
+        "gen_ai.usage.output_tokens",
+        _get_value(usage, "outputTokens"),
+    )
+    _set_if_present(
+        span,
+        "gen_ai.usage.total_tokens",
+        _get_value(usage, "totalTokens"),
+    )
+
+
+def _call_user_hook(hook: Optional[Callable[..., Any]], *args: Any) -> None:
+    if not callable(hook):
+        return
+    hook(*args)
+
+
 class BotocoreInstrumentationWrapper(BotocoreInstrumentor, BaseInstrumentorWrapper):
     '''Hypertrace wrapper around OTel Botocore Instrumentor class'''
+
+    def _instrument(self, **kwargs):
+        user_request_hook = kwargs.get("request_hook")
+        user_response_hook = kwargs.get("response_hook")
+
+        def request_hook(span, service_name, operation_name, api_params):
+            _call_user_hook(user_request_hook, span, service_name, operation_name, api_params)
+            if _is_bedrock_converse(service_name, operation_name):
+                _set_bedrock_request_attributes(span, operation_name, api_params)
+
+        def response_hook(span, service_name, operation_name, result):
+            _call_user_hook(user_response_hook, span, service_name, operation_name, result)
+            if _is_bedrock_converse(service_name, operation_name):
+                request_model = span.attributes.get("gen_ai.request.model")
+                _set_bedrock_response_attributes(span, result, request_model)
+
+        kwargs["request_hook"] = request_hook
+        kwargs["response_hook"] = response_hook
+        super()._instrument(**kwargs)

test-code/bedrock_direct_spans.py

@@ -0,0 +1,197 @@
+#!/usr/bin/env python3
+"""Call Bedrock Runtime directly with boto3 and write captured OTel spans to JSON."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import shlex
+import sys
+from pathlib import Path
+from typing import Any
+
+
+ROOT = Path(__file__).resolve().parents[1]
+WORKSPACE_ROOT = Path(__file__).resolve().parents[3]
+sys.path.insert(0, str(ROOT / "src"))
+
+
+def _load_env_file(path: Path) -> None:
+    if not path.exists():
+        raise FileNotFoundError(f"env file not found: {path}")
+    for raw_line in path.read_text(encoding="utf-8").splitlines():
+        line = raw_line.strip()
+        if not line or line.startswith("#"):
+            continue
+        if line.startswith("export "):
+            line = line[len("export ") :].strip()
+        try:
+            parts = shlex.split(line, posix=True)
+        except ValueError:
+            continue
+        if not parts or "=" not in parts[0]:
+            continue
+        key, value = parts[0].split("=", 1)
+        os.environ[key] = value
+
+
+def _default_env_path() -> Path:
+    workspace_env = WORKSPACE_ROOT / ".work" / ".env"
+    if workspace_env.exists():
+        return workspace_env
+    return ROOT / ".env"
+
+
+def _resolve_model_id(explicit_model_id: str | None) -> str:
+    if explicit_model_id:
+        return explicit_model_id
+    for key in ("BEDROCK_MODEL_ID", "BEDROCK_MODEL_ARN", "AWS_BEDROCK_MODEL_ID"):
+        value = os.environ.get(key)
+        if value:
+            return value
+    model_map = os.environ.get("BEDROCK_MODEL_MAP")
+    if model_map:
+        parsed = json.loads(model_map)
+        first_model = next(iter(parsed.values()))
+        if isinstance(first_model, dict):
+            arn = first_model.get("arn")
+            if arn:
+                return arn
+    raise ValueError("No Bedrock model found; pass --model-id or set BEDROCK_MODEL_MAP")
+
+
+def _response_text(response: dict[str, Any]) -> str:
+    content = (
+        response.get("output", {})
+        .get("message", {})
+        .get("content", [])
+    )
+    return "".join(part.get("text", "") for part in content if isinstance(part, dict))
+
+
+def _jsonable(value: Any) -> Any:
+    if value is None or isinstance(value, (bool, int, float, str)):
+        return value
+    if isinstance(value, (list, tuple)):
+        return [_jsonable(item) for item in value]
+    if isinstance(value, dict):
+        return {str(key): _jsonable(item) for key, item in value.items()}
+    return str(value)
+
+
+def _span_to_dict(span: Any) -> dict[str, Any]:
+    context = span.get_span_context()
+    parent = span.parent
+    return {
+        "name": span.name,
+        "context": {
+            "trace_id": format(context.trace_id, "032x"),
+            "span_id": format(context.span_id, "016x"),
+        },
+        "parent_span_id": format(parent.span_id, "016x") if parent else None,
+        "start_time": span.start_time,
+        "end_time": span.end_time,
+        "status": {
+            "status_code": str(span.status.status_code),
+            "description": span.status.description,
+        },
+        "attributes": _jsonable(dict(span.attributes or {})),
+    }
+
+
+class _StaticBearerToken:
+    def __init__(self, token: str) -> None:
+        self._token = token
+
+    def get_frozen_token(self) -> Any:
+        from botocore.tokens import FrozenAuthToken  # pylint: disable=import-outside-toplevel
+
+        return FrozenAuthToken(token=self._token)
+
+
+def _configure_bedrock_bearer_auth(client: Any) -> None:
+    token = os.environ.get("AWS_BEARER_TOKEN_BEDROCK")
+    if not token:
+        return
+    client._request_signer._auth_token = _StaticBearerToken(token)  # pylint: disable=protected-access
+    client.meta.events.register("choose-signer.bedrock-runtime", lambda **_kwargs: "bearer")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--env-file", type=Path, default=_default_env_path())
+    parser.add_argument("--spans-file", type=Path, default=ROOT / "test-code" / "bedrock_spans.json")
+    parser.add_argument("--model-id")
+    parser.add_argument("--prompt", default="Reply with one short sentence about OpenTelemetry.")
+    parser.add_argument("--max-tokens", type=int, default=64)
+    args = parser.parse_args()
+
+    _load_env_file(args.env_file)
+    os.environ.setdefault("HA_GEN_AI_ENABLED", "true")
+    os.environ.setdefault("HA_GEN_AI_PAYLOAD_CAPTURE_ENABLED", "false")
+    os.environ.setdefault("HA_GEN_AI_PAYLOAD_EVALUATION_ENABLED", "false")
+
+    import boto3  # pylint: disable=import-outside-toplevel
+    from opentelemetry.sdk.trace.export import SimpleSpanProcessor  # pylint: disable=import-outside-toplevel
+    from opentelemetry.sdk.trace.export.in_memory_span_exporter import (  # pylint: disable=import-outside-toplevel
+        InMemorySpanExporter,
+    )
+
+    from harness_sdk.agent import Agent  # pylint: disable=import-outside-toplevel
+
+    model_id = _resolve_model_id(args.model_id)
+    region = os.environ.get("AWS_REGION") or os.environ.get("AWS_DEFAULT_REGION")
+    if not region:
+        raise ValueError("AWS_REGION or AWS_DEFAULT_REGION is required")
+
+    agent = Agent()
+    exporter = InMemorySpanExporter()
+    agent.register_processor(SimpleSpanProcessor(exporter))
+    agent.instrument(
+        skip_libraries=[
+            "flask",
+            "django",
+            "fastapi",
+            "grpc:server",
+            "grpc:client",
+            "postgresql",
+            "mysql",
+            "requests",
+            "httpx",
+            "aiohttp:client",
+            "anthropic",
+            "openai",
+            "litellm",
+            "mcp",
+        ]
+    )
+
+    client = boto3.client("bedrock-runtime", region_name=region)
+    _configure_bedrock_bearer_auth(client)
+    response = client.converse(
+        modelId=model_id,
+        messages=[{"role": "user", "content": [{"text": args.prompt}]}],
+        inferenceConfig={"maxTokens": args.max_tokens},
+    )
+
+    spans = [_span_to_dict(span) for span in exporter.get_finished_spans()]
+    args.spans_file.parent.mkdir(parents=True, exist_ok=True)
+    args.spans_file.write_text(
+        json.dumps(
+            {
+                "model_id": model_id,
+                "region": region,
+                "response_text": _response_text(response),
+                "spans": spans,
+            },
+            indent=2,
+            sort_keys=True,
+        ),
+        encoding="utf-8",
+    )
+    print(f"Wrote {len(spans)} spans to {args.spans_file}")
+
+
+if __name__ == "__main__":
+    main()