telemetry-dev-bedrock 0.1.0__py3-none-any.whl

telemetry_dev_bedrock/__init__.py

@@ -0,0 +1,7 @@
+from __future__ import annotations
+
+from ._instrument import instrument_bedrock, uninstrument_bedrock, wrap_bedrock
+
+__version__ = "0.1.0"
+
+__all__ = ["__version__", "instrument_bedrock", "uninstrument_bedrock", "wrap_bedrock"]

telemetry_dev_bedrock/_fields.py

@@ -0,0 +1,102 @@
+from __future__ import annotations
+
+import json
+from typing import Any
+
+PROVIDER = "amazon-bedrock"
+
+
+def clean(fields: dict[str, Any]) -> dict[str, Any]:
+    return {key: value for key, value in fields.items() if value is not None}
+
+
+def metadata_fields(response: dict[str, Any] | BaseException | None) -> dict[str, Any]:
+    metadata: dict[str, Any] | None = None
+    if isinstance(response, dict) and isinstance(response.get("ResponseMetadata"), dict):
+        metadata = response["ResponseMetadata"]
+    else:
+        error_response = getattr(response, "response", None)
+        if isinstance(error_response, dict) and isinstance(
+            error_response.get("ResponseMetadata"), dict
+        ):
+            metadata = error_response["ResponseMetadata"]
+    if metadata is None:
+        return {}
+    retry_attempts = metadata.get("RetryAttempts")
+    attempts = retry_attempts + 1 if isinstance(retry_attempts, int) else None
+    total_retry_delay = metadata.get("TotalRetryDelay")
+    if total_retry_delay is None:
+        total_retry_delay = metadata.get("totalRetryDelay")
+    attributes = clean(
+        {
+            "aws.http.status_code": metadata.get("HTTPStatusCode"),
+            "aws.request.attempts": attempts if attempts and attempts > 1 else None,
+            "aws.request.total_retry_delay_ms": total_retry_delay,
+        }
+    )
+    return clean(
+        {
+            "response_id": metadata.get("RequestId"),
+            "attributes": attributes or None,
+        }
+    )
+
+
+def error_fields(error: BaseException) -> dict[str, Any]:
+    fields = metadata_fields(error)
+    response = getattr(error, "response", None)
+    code = None
+    if isinstance(response, dict) and isinstance(response.get("Error"), dict):
+        code = response["Error"].get("Code")
+    fields["error"] = error
+    if code:
+        attrs = dict(fields.get("attributes") or {})
+        attrs["aws.error.code"] = code
+        fields["attributes"] = attrs
+    return fields
+
+
+def usage_from_converse(value: Any) -> dict[str, int] | None:
+    if not isinstance(value, dict):
+        return None
+    usage = clean(
+        {
+            "input_tokens": value.get("inputTokens"),
+            "output_tokens": value.get("outputTokens"),
+            "total_tokens": value.get("totalTokens"),
+            "cache_read_input_tokens": value.get("cacheReadInputTokens"),
+            "cache_creation_input_tokens": value.get("cacheWriteInputTokens"),
+        }
+    )
+    return usage or None
+
+
+def metadata_attr_value(value: Any) -> Any:
+    if value is None or isinstance(value, str):
+        return value
+    return json.dumps(value, default=repr, ensure_ascii=False)
+
+
+def merge_fields(*parts: dict[str, Any]) -> dict[str, Any]:
+    merged: dict[str, Any] = {}
+    for part in parts:
+        metadata = part.get("metadata")
+        attributes = part.get("attributes")
+        usage = part.get("usage")
+        for key, value in part.items():
+            if key not in {"metadata", "attributes", "usage"} and value is not None:
+                merged[key] = value
+        if isinstance(metadata, dict):
+            merged["metadata"] = {
+                **merged.get("metadata", {}),
+                **{
+                    key: normalized
+                    for key, value in metadata.items()
+                    if (normalized := metadata_attr_value(value)) is not None
+                },
+            }
+        if isinstance(attributes, dict):
+            merged["attributes"] = {**merged.get("attributes", {}), **attributes}
+        if isinstance(usage, dict):
+            merged["usage"] = {**merged.get("usage", {}), **usage}
+    return merged

telemetry_dev_bedrock/_instrument.py

@@ -0,0 +1,231 @@
+from __future__ import annotations
+
+import io
+import threading
+import time
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any, TypeVar, cast
+
+import telemetry_dev
+from botocore.client import BaseClient
+from botocore.exceptions import ClientError
+from botocore.response import StreamingBody
+
+from ._fields import error_fields, merge_fields, metadata_fields
+from ._invoke_model import parse_body
+from ._registry import OperationSpec, lookup_operation
+from ._streams import InstrumentedEventStream
+
+_T = TypeVar("_T", bound=BaseClient)
+_SUPPORTED_SERVICES = frozenset({"bedrock-runtime", "bedrock-agent-runtime"})
+_WRAPPED_ATTR = "_telemetry_dev_bedrock_wrapped"
+
+
+@dataclass(frozen=True)
+class _Options:
+    capture_agent_trace: bool = False
+
+
+_original_api_call: Callable[..., Any] | None = None
+_patched_api_call: Callable[..., Any] | None = None
+_LOCK = threading.Lock()
+
+
+def wrap_bedrock(client: _T, *, capture_agent_trace: bool = False) -> _T:
+    with _LOCK:
+        if getattr(client, _WRAPPED_ATTR, False):
+            return client
+        options = _Options(capture_agent_trace=capture_agent_trace)
+        mapping = getattr(client.meta, "method_to_api_mapping", {})
+        if not isinstance(mapping, dict):
+            return client
+        for method_name, operation_name in mapping.items():
+            if not isinstance(method_name, str) or not isinstance(operation_name, str):
+                continue
+            service_name = client.meta.service_model.service_name
+            if lookup_operation(service_name, operation_name) is None or not hasattr(
+                client, method_name
+            ):
+                continue
+            original = getattr(client, method_name)
+            setattr(
+                client, method_name, _make_method_wrapper(client, operation_name, original, options)
+            )
+        setattr(client, _WRAPPED_ATTR, True)
+        return client
+
+
+def instrument_bedrock(*, capture_agent_trace: bool = False) -> None:
+    options = _Options(capture_agent_trace=capture_agent_trace)
+    with _LOCK:
+        global _original_api_call, _patched_api_call
+        if _original_api_call is not None:
+            return
+        original = BaseClient._make_api_call
+        _original_api_call = original
+
+        def patched(self: BaseClient, operation_name: str, api_params: dict[str, Any]) -> Any:
+            if getattr(self, _WRAPPED_ATTR, False):
+                return original(self, operation_name, api_params)
+            return _call_with_span(
+                lambda op_name, params: original(self, op_name, params),
+                self,
+                operation_name,
+                api_params,
+                options,
+            )
+
+        _patched_api_call = patched
+        BaseClient._make_api_call = patched  # type: ignore[method-assign]
+
+
+def uninstrument_bedrock() -> None:
+    with _LOCK:
+        global _original_api_call, _patched_api_call
+        if _original_api_call is None or _patched_api_call is None:
+            return
+        if BaseClient._make_api_call is not _patched_api_call:
+            return
+        BaseClient._make_api_call = _original_api_call  # type: ignore[method-assign]
+        _original_api_call = None
+        _patched_api_call = None
+
+
+def _make_method_wrapper(
+    client: BaseClient,
+    operation_name: str,
+    original: Callable[..., Any],
+    options: _Options,
+) -> Callable[..., Any]:
+    def wrapped(*args: Any, **kwargs: Any) -> Any:
+        if args:
+            return original(*args, **kwargs)
+        return _call_with_span(
+            lambda _op, params: original(**params), client, operation_name, kwargs, options
+        )
+
+    return wrapped
+
+
+def _call_with_span(
+    original: Callable[[str, dict[str, Any]], Any],
+    client: BaseClient,
+    operation_name: str,
+    api_params: dict[str, Any],
+    options: _Options,
+) -> Any:
+    service_name = client.meta.service_model.service_name
+    if service_name not in _SUPPORTED_SERVICES:
+        return original(operation_name, api_params)
+    spec = lookup_operation(service_name, operation_name)
+    if spec is None:
+        return original(operation_name, api_params)
+
+    started_at = time.perf_counter()
+    name, start_fields = _safe_start_fields(spec, api_params)
+    handle = telemetry_dev.start_span(name, **start_fields)
+    end = _end_once(handle)
+
+    try:
+        response = original(operation_name, api_params)
+    except ClientError as exc:
+        end(**error_fields(exc))
+        raise
+    except BaseException as exc:
+        end(error=exc)
+        raise
+
+    if not isinstance(response, dict):
+        end()
+        return response
+
+    if operation_name == "InvokeModel" and isinstance(response.get("body"), StreamingBody):
+        return _handle_streaming_body(spec, api_params, response, end)
+
+    stream_key = spec.stream_key
+    if stream_key is not None and response.get(stream_key) is not None:
+        state_factory = spec.stream_state_factory
+        if state_factory is None:
+            end(**metadata_fields(response))
+            return response
+        state = state_factory(options.capture_agent_trace)
+        try:
+            base_fields = merge_fields(
+                metadata_fields(response), spec.response_fields(api_params, response)
+            )
+        except Exception:
+            base_fields = metadata_fields(response)
+
+        def finish(fields: dict[str, Any]) -> None:
+            if fields.keys() == {"time_to_first_chunk_ms"}:
+                try:
+                    handle.update(**fields)
+                except Exception:
+                    pass
+                return
+            end(**merge_fields(base_fields, fields))
+
+        return {
+            **response,
+            stream_key: InstrumentedEventStream(response[stream_key], state, finish, started_at),
+        }
+
+    end(**merge_fields(metadata_fields(response), spec.response_fields(api_params, response)))
+    return response
+
+
+def _handle_streaming_body(
+    spec: OperationSpec,
+    params: dict[str, Any],
+    response: dict[str, Any],
+    end: Callable[..., None],
+) -> dict[str, Any]:
+    body = cast(StreamingBody, response["body"])
+    try:
+        raw = body.read()
+    except Exception:
+        end(**metadata_fields(response))
+        return response
+    parsed = parse_body(raw, response.get("contentType"))
+    response_for_fields = {**response, "_telemetry_dev_parsed_body": parsed}
+    try:
+        fields = merge_fields(
+            metadata_fields(response), spec.response_fields(params, response_for_fields)
+        )
+    except Exception:
+        fields = metadata_fields(response)
+    end(**fields)
+    return {**response, "body": StreamingBody(io.BytesIO(raw), len(raw))}
+
+
+def _safe_start_fields(spec: OperationSpec, params: dict[str, Any]) -> tuple[str, dict[str, Any]]:
+    try:
+        name = spec.span_name(params)
+    except Exception:
+        name = "bedrock"
+    try:
+        fields = spec.request_fields(params)
+    except Exception:
+        fields = {}
+    try:
+        fields = {"type": spec.span_type(params), **fields}
+    except Exception:
+        fields = {"type": "span", **fields}
+    return name, fields
+
+
+def _end_once(handle: telemetry_dev.SpanHandle) -> Callable[..., None]:
+    ended = False
+
+    def end(**fields: Any) -> None:
+        nonlocal ended
+        if ended:
+            return
+        ended = True
+        try:
+            handle.end(**fields)
+        except Exception:
+            pass
+
+    return end

telemetry_dev_bedrock/_invoke_model.py

@@ -0,0 +1,241 @@
+from __future__ import annotations
+
+import json
+from typing import Any
+
+from ._fields import clean
+
+
+def parse_body(body: Any, content_type: Any = "application/json") -> Any:
+    if isinstance(content_type, str) and "json" not in content_type:
+        return None
+    if isinstance(body, bytes | bytearray):
+        raw = bytes(body).decode("utf-8")
+    elif isinstance(body, str):
+        raw = body
+    elif all(hasattr(body, name) for name in ("tell", "read", "seek")):
+        seekable = getattr(body, "seekable", None)
+        if callable(seekable):
+            try:
+                if not seekable():
+                    return None
+            except Exception:
+                return None
+        try:
+            position = body.tell()
+            body.seek(position)
+            value = body.read()
+            body.seek(position)
+        except Exception:
+            return None
+        if isinstance(value, bytes | bytearray):
+            raw = bytes(value).decode("utf-8")
+        elif isinstance(value, str):
+            raw = value
+        else:
+            return None
+    else:
+        return None
+    try:
+        return json.loads(raw)
+    except Exception:
+        return None
+
+
+def is_embedding_model(model_id: Any) -> bool:
+    model = model_id.lower() if isinstance(model_id, str) else ""
+    return "titan-embed" in model or "cohere.embed" in model or "-embedding" in model
+
+
+def request_fields(params: dict[str, Any]) -> dict[str, Any]:
+    model = params.get("modelId") if isinstance(params.get("modelId"), str) else None
+    body = parse_body(params.get("body"), params.get("contentType"))
+    fields = {"provider": "amazon-bedrock", "model": model, "input": body}
+    if is_embedding_model(model):
+        fields["output_type"] = "embedding"
+    fields.update(sampling_fields(model, body))
+    return clean(fields)
+
+
+def sampling_fields(model: str | None, body: Any) -> dict[str, Any]:
+    if not isinstance(body, dict):
+        return {}
+    lower = model.lower() if isinstance(model, str) else ""
+    if "amazon.titan" in lower:
+        cfg = (
+            body.get("textGenerationConfig")
+            if isinstance(body.get("textGenerationConfig"), dict)
+            else {}
+        )
+        return clean(
+            {
+                "temperature": cfg.get("temperature"),
+                "top_p": cfg.get("topP"),
+                "max_tokens": cfg.get("maxTokenCount"),
+                "stop_sequences": cfg.get("stopSequences"),
+            }
+        )
+    if "amazon.nova" in lower:
+        cfg = body.get("inferenceConfig") if isinstance(body.get("inferenceConfig"), dict) else {}
+        return clean(
+            {
+                "temperature": cfg.get("temperature"),
+                "top_p": cfg.get("topP") or cfg.get("top_p"),
+                "top_k": cfg.get("topK"),
+                "max_tokens": cfg.get("maxTokens") or cfg.get("max_new_tokens"),
+                "stop_sequences": cfg.get("stopSequences"),
+            }
+        )
+    if "anthropic.claude" in lower:
+        return clean(
+            {
+                "temperature": body.get("temperature"),
+                "top_p": body.get("top_p"),
+                "top_k": body.get("top_k"),
+                "max_tokens": body.get("max_tokens"),
+                "stop_sequences": body.get("stop_sequences"),
+            }
+        )
+    if "meta.llama" in lower:
+        return clean(
+            {
+                "temperature": body.get("temperature"),
+                "top_p": body.get("top_p"),
+                "max_tokens": body.get("max_gen_len"),
+            }
+        )
+    return clean(
+        {
+            "temperature": body.get("temperature"),
+            "top_p": body.get("top_p") or body.get("topP"),
+            "max_tokens": body.get("max_tokens") or body.get("maxTokens"),
+        }
+    )
+
+
+def response_fields(
+    model: str | None, body: Any, headers: dict[str, Any] | None = None
+) -> dict[str, Any]:
+    if not isinstance(body, dict):
+        return header_usage_fields(headers)
+    lower = model.lower() if isinstance(model, str) else ""
+    if is_embedding_model(model):
+        return _with_header_usage(
+            {
+                "output": body,
+                "output_type": "embedding",
+                "usage": clean({"input_tokens": body.get("inputTextTokenCount")}) or None,
+            },
+            headers,
+        )
+    if "amazon.titan" in lower:
+        first = body.get("results", [{}])[0] if isinstance(body.get("results"), list) else {}
+        first = first if isinstance(first, dict) else {}
+        return clean(
+            {
+                "output": body,
+                "finish_reason": first.get("completionReason"),
+                "usage": clean(
+                    {
+                        "input_tokens": body.get("inputTextTokenCount"),
+                        "output_tokens": first.get("tokenCount"),
+                    }
+                )
+                or None,
+            }
+        )
+    if "amazon.nova" in lower:
+        usage = body.get("usage") if isinstance(body.get("usage"), dict) else {}
+        return clean(
+            {
+                "output": body,
+                "finish_reason": body.get("stopReason"),
+                "usage": clean(
+                    {
+                        "input_tokens": usage.get("inputTokens"),
+                        "output_tokens": usage.get("outputTokens"),
+                        "total_tokens": usage.get("totalTokens"),
+                    }
+                )
+                or None,
+            }
+        )
+    if "anthropic.claude" in lower:
+        usage = body.get("usage") if isinstance(body.get("usage"), dict) else {}
+        parsed_usage = clean(
+            {
+                "input_tokens": usage.get("input_tokens"),
+                "output_tokens": usage.get("output_tokens"),
+                "cache_read_input_tokens": usage.get("cache_read_input_tokens"),
+                "cache_creation_input_tokens": usage.get("cache_creation_input_tokens"),
+            }
+        )
+        return clean(
+            {
+                "output": body,
+                "finish_reason": body.get("stop_reason"),
+                "usage": parsed_usage or header_usage_fields(headers).get("usage"),
+            }
+        )
+    if "meta.llama" in lower:
+        return clean(
+            {
+                "output": body,
+                "finish_reason": body.get("stop_reason"),
+                "usage": clean(
+                    {
+                        "input_tokens": body.get("prompt_token_count"),
+                        "output_tokens": body.get("generation_token_count"),
+                    }
+                )
+                or None,
+            }
+        )
+    if "mistral" in lower:
+        first = body.get("outputs", [{}])[0] if isinstance(body.get("outputs"), list) else {}
+        first = first if isinstance(first, dict) else {}
+        return _with_header_usage(
+            {"output": body, "finish_reason": first.get("stop_reason")}, headers
+        )
+    if "cohere.command-r" in lower:
+        return _with_header_usage(
+            {"output": body, "finish_reason": body.get("finish_reason")}, headers
+        )
+    if "cohere.command" in lower:
+        first = (
+            body.get("generations", [{}])[0] if isinstance(body.get("generations"), list) else {}
+        )
+        first = first if isinstance(first, dict) else {}
+        return _with_header_usage(
+            {"output": body, "finish_reason": first.get("finish_reason")}, headers
+        )
+    return _with_header_usage({"output": body}, headers)
+
+
+def _with_header_usage(fields: dict[str, Any], headers: dict[str, Any] | None) -> dict[str, Any]:
+    cleaned = clean(fields)
+    if cleaned.get("usage"):
+        return cleaned
+    return clean({**cleaned, **header_usage_fields(headers)})
+
+
+def header_usage_fields(headers: dict[str, Any] | None) -> dict[str, Any]:
+    if not headers:
+        return {}
+    input_tokens = _int_header(headers, "x-amzn-bedrock-input-token-count")
+    output_tokens = _int_header(headers, "x-amzn-bedrock-output-token-count")
+    return clean(
+        {
+            "usage": clean({"input_tokens": input_tokens, "output_tokens": output_tokens}) or None,
+        }
+    )
+
+
+def _int_header(headers: dict[str, Any], key: str) -> int | None:
+    value = headers.get(key)
+    if value is None:
+        return None
+    try:
+        return int(value)
+    except (TypeError, ValueError):
+        return None