diff --git a/src/harness_sdk/instrumentation/botocore/__init__.py b/src/harness_sdk/instrumentation/botocore/__init__.py index f5f796c..525cb1c 100644 --- a/src/harness_sdk/instrumentation/botocore/__init__.py +++ b/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) diff --git a/test-code/bedrock_direct_spans.py b/test-code/bedrock_direct_spans.py new file mode 100644 index 0000000..0e601cf --- /dev/null +++ b/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() diff --git a/test/instrumentation/botocore/test_bedrock_genai.py b/test/instrumentation/botocore/test_bedrock_genai.py new file mode 100644 index 0000000..9677880 --- /dev/null +++ b/test/instrumentation/botocore/test_bedrock_genai.py @@ -0,0 +1,168 @@ +import boto3 +import pytest +from botocore.stub import Stubber + +from harness_sdk.instrumentation.botocore import BotocoreInstrumentationWrapper + + +_MODEL_ID = "anthropic.claude-3-haiku-20240307-v1:0" +_INFERENCE_PROFILE_ARN = ( + "arn:aws:bedrock:us-east-1:123456789012:" + "application-inference-profile/profile123" +) + + +@pytest.fixture +def botocore_instrumentor(): + wrapper = BotocoreInstrumentationWrapper() + yield wrapper + if getattr(wrapper, "_is_instrumented_by_opentelemetry", False): + wrapper.uninstrument() + + +def _bedrock_span(spans): + for span in spans: + attrs = span.attributes or {} + if attrs.get("gen_ai.provider.name") == "aws.bedrock": + return span + raise AssertionError("No Bedrock GenAI span found") + + +def _converse_params(model_id=_MODEL_ID): + return { + "modelId": model_id, + "messages": [{"role": "user", "content": [{"text": "hi"}]}], + "inferenceConfig": { + "maxTokens": 32, + "temperature": 0.1, + "topP": 0.9, + }, + } + + +def test_bedrock_converse_span_has_gen_ai_attributes(agent, exporter, botocore_instrumentor): # pylint: disable=unused-argument + botocore_instrumentor.instrument() + client = boto3.client( + "bedrock-runtime", + region_name="us-east-1", + aws_access_key_id="test", + aws_secret_access_key="test", + ) + response = { + "output": { + "message": { + "role": "assistant", + "content": [{"text": "hello"}], + } + }, + "stopReason": "end_turn", + "usage": {"inputTokens": 3, "outputTokens": 5, "totalTokens": 8}, + "metrics": {"latencyMs": 42}, + "ResponseMetadata": { + "HTTPStatusCode": 200, + "HTTPHeaders": {"x-amzn-bedrock-model-id": _MODEL_ID}, + }, + } + + with Stubber(client) as stubber: + stubber.add_response("converse", response, _converse_params(_INFERENCE_PROFILE_ARN)) + client.converse(**_converse_params(_INFERENCE_PROFILE_ARN)) + + attrs = _bedrock_span(exporter.get_finished_spans()).attributes + exporter.clear() + assert attrs.get("gen_ai.request.model") == _INFERENCE_PROFILE_ARN + assert attrs.get("gen_ai.operation.name") == "chat" + assert attrs.get("gen_ai.provider.name") == "aws.bedrock" + assert attrs.get("gen_ai.framework") == "boto3" + assert attrs.get("aws.bedrock.inference_profile_arn") == _INFERENCE_PROFILE_ARN + assert attrs.get("gen_ai.request.streaming") is False + assert attrs.get("gen_ai.request.max_tokens") == 32 + assert attrs.get("gen_ai.request.temperature") == 0.1 + assert attrs.get("gen_ai.request.top_p") == 0.9 + assert attrs.get("gen_ai.response.model") == _MODEL_ID + assert attrs.get("aws.bedrock.execution_model_id") == _MODEL_ID + assert attrs.get("gen_ai.response.finish_reasons") == "end_turn" + assert attrs.get("gen_ai.usage.input_tokens") == 3 + assert attrs.get("gen_ai.usage.output_tokens") == 5 + assert attrs.get("gen_ai.usage.total_tokens") == 8 + + +def test_bedrock_converse_falls_back_to_request_model_when_response_header_missing( + agent, + exporter, + botocore_instrumentor, +): # pylint: disable=unused-argument + botocore_instrumentor.instrument() + client = boto3.client( + "bedrock-runtime", + region_name="us-east-1", + aws_access_key_id="test", + aws_secret_access_key="test", + ) + response = { + "output": { + "message": { + "role": "assistant", + "content": [{"text": "hello"}], + } + }, + "stopReason": "end_turn", + "usage": {"inputTokens": 3, "outputTokens": 5, "totalTokens": 8}, + "metrics": {"latencyMs": 42}, + "ResponseMetadata": {"HTTPStatusCode": 200, "HTTPHeaders": {}}, + } + + with Stubber(client) as stubber: + stubber.add_response("converse", response, _converse_params(_INFERENCE_PROFILE_ARN)) + client.converse(**_converse_params(_INFERENCE_PROFILE_ARN)) + + attrs = _bedrock_span(exporter.get_finished_spans()).attributes + exporter.clear() + assert attrs.get("gen_ai.request.model") == _INFERENCE_PROFILE_ARN + assert attrs.get("gen_ai.response.model") == _INFERENCE_PROFILE_ARN + assert attrs.get("aws.bedrock.execution_model_id") is None + assert attrs.get("aws.bedrock.inference_profile_arn") == _INFERENCE_PROFILE_ARN + assert "aws.bedrock.inference_profile_model_arns" not in attrs + assert "aws.bedrock.resolved_model_id" not in attrs + + +def test_bedrock_converse_stream_request_has_gen_ai_attributes(agent, exporter, botocore_instrumentor): # pylint: disable=unused-argument + calls = [] + + def request_hook(span, service_name, operation_name, api_params): + calls.append((service_name, operation_name, api_params["modelId"])) + + botocore_instrumentor.instrument(request_hook=request_hook) + + client = boto3.client( + "bedrock-runtime", + region_name="us-east-1", + aws_access_key_id="test", + aws_secret_access_key="test", + ) + params = { + "modelId": _MODEL_ID, + "messages": [{"role": "user", "content": [{"text": "hi"}]}], + } + + with Stubber(client) as stubber: + stubber.add_client_error( + "converse_stream", + service_error_code="ThrottlingException", + service_message="rate limited", + http_status_code=429, + expected_params=params, + ) + try: + client.converse_stream(**params) + except client.exceptions.ThrottlingException: + pass + + attrs = _bedrock_span(exporter.get_finished_spans()).attributes + exporter.clear() + assert calls == [("bedrock-runtime", "ConverseStream", _MODEL_ID)] + assert attrs.get("gen_ai.request.model") == _MODEL_ID + assert attrs.get("gen_ai.operation.name") == "chat" + assert attrs.get("gen_ai.provider.name") == "aws.bedrock" + assert attrs.get("gen_ai.framework") == "boto3" + assert attrs.get("gen_ai.request.streaming") is True