frugal-sdk-python 0.1.0__py3-none-any.whl

frugal_metrics/wrappers/anthropic.py

@@ -0,0 +1,277 @@
+"""Anthropic client wrapper.
+
+Monkey-patches `messages.create` on an Anthropic / AsyncAnthropic /
+AnthropicBedrock / AnthropicVertex client instance. Handles non-streaming
+responses and `stream=True` streams. The `messages.stream(...)` context-manager
+API is not yet wrapped (deferred).
+"""
+
+from __future__ import annotations
+
+import functools
+import inspect
+import logging
+from typing import Any, Callable, TypeVar
+
+from frugal_metrics import caller as caller_mod
+from frugal_metrics import semconv
+from frugal_metrics.config import load_config
+from frugal_metrics.prompt_analyzer import RequestPayload
+from frugal_metrics.safe import safe_call
+from frugal_metrics.instrument import (
+    CallContext,
+    UsageReport,
+    start_call,
+)
+from frugal_metrics.streams import AsyncStreamProxy, SyncStreamProxy
+
+logger = logging.getLogger("frugal_metrics")
+
+T = TypeVar("T")
+
+
+def wrap_anthropic(client: T) -> T:
+    """Instrument an Anthropic (or Bedrock / Vertex variant) client in place."""
+    config = load_config()
+    if config is None:
+        return client
+
+    try:
+        import anthropic  # type: ignore
+    except ImportError:
+        logger.warning(
+            "frugal-metrics: anthropic package not installed; cannot wrap"
+        )
+        return client
+
+    if getattr(client, "_frugal_wrapped", False):
+        return client
+
+    system = _detect_system(client)
+    skip_dirs = _build_skip_dirs(anthropic)
+
+    messages = getattr(client, "messages", None)
+    if messages is None:
+        logger.debug(
+            "frugal-metrics: no `messages` attribute on %s; skipping",
+            type(client).__name__,
+        )
+        return client
+
+    original = getattr(messages, "create", None)
+    if original is None or getattr(original, "_frugal_wrapped", False):
+        return client
+
+    is_async = inspect.iscoroutinefunction(original)
+    wrapped = (
+        _make_async(original, system, skip_dirs, config)
+        if is_async
+        else _make_sync(original, system, skip_dirs, config)
+    )
+    wrapped._frugal_wrapped = True  # type: ignore[attr-defined]
+    try:
+        setattr(messages, "create", wrapped)
+    except Exception as exc:  # noqa: BLE001
+        logger.warning("frugal-metrics: failed to patch messages.create: %s", exc)
+        return client
+
+    try:
+        setattr(client, "_frugal_wrapped", True)
+    except Exception:  # noqa: BLE001
+        pass
+    return client
+
+
+def _detect_system(client: Any) -> str:
+    cls_name = type(client).__name__
+    if "Bedrock" in cls_name:
+        return semconv.SYSTEM_AWS_BEDROCK
+    if "Vertex" in cls_name:
+        return semconv.SYSTEM_VERTEX_AI
+    return semconv.SYSTEM_ANTHROPIC
+
+
+def _build_skip_dirs(anthropic_mod: Any) -> tuple[str, ...]:
+    import frugal_metrics
+
+    return tuple(
+        d
+        for d in (
+            caller_mod.package_dir_for(frugal_metrics),
+            caller_mod.package_dir_for(anthropic_mod),
+        )
+        if d
+    )
+
+
+def _make_sync(
+    original: Callable[..., Any],
+    system: str,
+    skip_dirs: tuple[str, ...],
+    config: Any,
+) -> Callable[..., Any]:
+    @functools.wraps(original)
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
+        try:
+            streaming = bool(kwargs.get("stream"))
+            ctx = CallContext(
+                system=system,
+                operation=semconv.OP_CHAT,
+                request_model=_request_model(kwargs),
+                skip_package_dirs=skip_dirs,
+                batched=False,
+                request_payload=_capture_anthropic_payload(kwargs),
+            )
+            emit = start_call(config, ctx)
+        except Exception as exc:  # noqa: BLE001
+            _record_internal_error("setup_failed", exc, config)
+            return original(*args, **kwargs)
+
+        try:
+            response = original(*args, **kwargs)
+        except BaseException as exc:
+            safe_call(lambda: emit.error(exc))
+            raise
+
+        if streaming:
+            try:
+                return SyncStreamProxy(response, emit, _anthropic_stream_event)
+            except Exception as exc:  # noqa: BLE001
+                _record_internal_error("stream_proxy_failed", exc, config)
+                return response
+
+        safe_call(lambda: emit.success(_messages_usage(response)))
+        return response
+
+    return wrapper
+
+
+def _make_async(
+    original: Callable[..., Any],
+    system: str,
+    skip_dirs: tuple[str, ...],
+    config: Any,
+) -> Callable[..., Any]:
+    @functools.wraps(original)
+    async def wrapper(*args: Any, **kwargs: Any) -> Any:
+        try:
+            streaming = bool(kwargs.get("stream"))
+            ctx = CallContext(
+                system=system,
+                operation=semconv.OP_CHAT,
+                request_model=_request_model(kwargs),
+                skip_package_dirs=skip_dirs,
+                batched=False,
+                request_payload=_capture_anthropic_payload(kwargs),
+            )
+            emit = start_call(config, ctx)
+        except Exception as exc:  # noqa: BLE001
+            _record_internal_error("setup_failed", exc, config)
+            return await original(*args, **kwargs)
+
+        try:
+            response = await original(*args, **kwargs)
+        except BaseException as exc:
+            safe_call(lambda: emit.error(exc))
+            raise
+
+        if streaming:
+            try:
+                return AsyncStreamProxy(response, emit, _anthropic_stream_event)
+            except Exception as exc:  # noqa: BLE001
+                _record_internal_error("stream_proxy_failed", exc, config)
+                return response
+
+        safe_call(lambda: emit.success(_messages_usage(response)))
+        return response
+
+    return wrapper
+
+
+def _record_internal_error(stage: str, exc: BaseException, config: Any) -> None:
+    """Try to emit an internal error metric; swallow any failure."""
+    try:
+        from frugal_metrics.otel import get_instruments
+
+        instruments = get_instruments(config)
+        err_attrs: dict[str, Any] = {
+            semconv.FRUGAL_PROJECT_ID: config.project_id,
+            "frugal.internal_error.stage": stage,
+            "frugal.internal_error.type": type(exc).__name__,
+        }
+        if config.customer_id:
+            err_attrs[semconv.FRUGAL_CUSTOMER_ID] = config.customer_id
+        instruments.internal_errors.add(1, err_attrs)
+    except Exception:  # noqa: BLE001
+        pass
+
+
+def _request_model(kwargs: dict[str, Any]) -> str:
+    model = kwargs.get("model")
+    return str(model) if model is not None else "<unknown>"
+
+
+def _capture_anthropic_payload(kwargs: dict[str, Any]) -> RequestPayload:
+    max_tokens = kwargs.get("max_tokens")
+    if not isinstance(max_tokens, int):
+        max_tokens = None
+    return RequestPayload(
+        messages=kwargs.get("messages"),
+        tools=kwargs.get("tools"),
+        response_format=None,
+        max_tokens=max_tokens,
+        system=kwargs.get("system"),
+    )
+
+
+def _attr(obj: Any, name: str, default: Any = None) -> Any:
+    if obj is None:
+        return default
+    if isinstance(obj, dict):
+        return obj.get(name, default)
+    return getattr(obj, name, default)
+
+
+def _messages_usage(response: Any) -> UsageReport:
+    usage = _attr(response, "usage")
+    return UsageReport(
+        response_model=_attr(response, "model"),
+        input_tokens=_attr(usage, "input_tokens"),
+        output_tokens=_attr(usage, "output_tokens"),
+        cache_read_tokens=_attr(usage, "cache_read_input_tokens"),
+        cache_write_tokens=_attr(usage, "cache_creation_input_tokens"),
+    )
+
+
+def _anthropic_stream_event(event: Any, usage: UsageReport) -> None:
+    """Accumulate usage from Anthropic stream events.
+
+    Event types we care about:
+    - `message_start`: event.message.usage has initial input_tokens + cache_* and
+      a starting output_tokens value.
+    - `message_delta`: event.usage.output_tokens is the running output total.
+    """
+    event_type = _attr(event, "type")
+    if event_type == "message_start":
+        message = _attr(event, "message")
+        model = _attr(message, "model")
+        if model:
+            usage.response_model = model
+        m_usage = _attr(message, "usage")
+        if m_usage is not None:
+            usage.input_tokens = _attr(m_usage, "input_tokens", usage.input_tokens)
+            usage.output_tokens = _attr(
+                m_usage, "output_tokens", usage.output_tokens
+            )
+            cache_read = _attr(m_usage, "cache_read_input_tokens")
+            if cache_read is not None:
+                usage.cache_read_tokens = cache_read
+            cache_write = _attr(m_usage, "cache_creation_input_tokens")
+            if cache_write is not None:
+                usage.cache_write_tokens = cache_write
+    elif event_type == "message_delta":
+        d_usage = _attr(event, "usage")
+        if d_usage is not None:
+            out = _attr(d_usage, "output_tokens")
+            if out is not None:
+                usage.output_tokens = out

frugal_metrics/wrappers/bedrock.py

@@ -0,0 +1,358 @@
+"""boto3 / aiobotocore Bedrock client wrapper.
+
+Monkey-patches the four bedrock-runtime operations on the customer's client
+instance: `invoke_model`, `invoke_model_with_response_stream`, `converse`,
+`converse_stream`. Detects sync (boto3) vs async (aiobotocore) per-method by
+checking whether the bound method is a coroutine function — same pattern as
+`wrap_openai` for `OpenAI` vs `AsyncOpenAI`.
+
+Per-provider response-shape extraction lives in `bedrock_models.py`; this
+file is only the glue that hooks into the existing `instrument.start_call`
+emitter and the `SyncStreamProxy` / `AsyncStreamProxy` infrastructure.
+"""
+
+from __future__ import annotations
+
+import functools
+import inspect
+import logging
+from typing import Any, Callable, TypeVar
+
+from frugal_metrics import caller as caller_mod
+from frugal_metrics import semconv
+from frugal_metrics.config import load_config
+from frugal_metrics.instrument import (
+    CallContext,
+    UsageReport,
+    start_call,
+)
+from frugal_metrics.safe import safe_call
+from frugal_metrics.streams import AsyncStreamProxy, SyncStreamProxy
+from frugal_metrics.wrappers import bedrock_models
+
+logger = logging.getLogger("frugal_metrics")
+
+T = TypeVar("T")
+
+# The four bedrock-runtime operations we wrap.
+_OP_INVOKE_MODEL = "invoke_model"
+_OP_INVOKE_MODEL_STREAM = "invoke_model_with_response_stream"
+_OP_CONVERSE = "converse"
+_OP_CONVERSE_STREAM = "converse_stream"
+
+_STREAMING_OPS = frozenset({_OP_INVOKE_MODEL_STREAM, _OP_CONVERSE_STREAM})
+_INVOKE_MODEL_OPS = frozenset({_OP_INVOKE_MODEL, _OP_INVOKE_MODEL_STREAM})
+
+
+def wrap_bedrock(client: T) -> T:
+    """Instrument a boto3 / aiobotocore bedrock-runtime client in place.
+
+    Silently no-ops if:
+    - frugal-metrics config is missing,
+    - the client isn't actually bedrock-runtime (it's a different boto3 service),
+    - or the client is already wrapped.
+    """
+    config = load_config()
+    if config is None:
+        return client
+
+    if getattr(client, "_frugal_wrapped", False):
+        return client
+
+    if not _is_bedrock_runtime_client(client):
+        return client
+
+    skip_dirs = _build_skip_dirs(client)
+
+    for op_name in (
+        _OP_INVOKE_MODEL,
+        _OP_INVOKE_MODEL_STREAM,
+        _OP_CONVERSE,
+        _OP_CONVERSE_STREAM,
+    ):
+        original = getattr(client, op_name, None)
+        if original is None or getattr(original, "_frugal_wrapped", False):
+            continue
+        wrapped = (
+            _make_async(op_name, original, skip_dirs, config)
+            if inspect.iscoroutinefunction(original)
+            else _make_sync(op_name, original, skip_dirs, config)
+        )
+        try:
+            wrapped._frugal_wrapped = True  # type: ignore[attr-defined]
+            setattr(client, op_name, wrapped)
+        except Exception as exc:  # noqa: BLE001
+            logger.warning(
+                "frugal-metrics: failed to patch bedrock-runtime.%s: %s", op_name, exc
+            )
+
+    try:
+        setattr(client, "_frugal_wrapped", True)
+    except Exception:  # noqa: BLE001
+        pass
+    return client
+
+
+# --- Internal helpers --------------------------------------------------------
+
+
+def _is_bedrock_runtime_client(client: Any) -> bool:
+    meta = getattr(client, "meta", None)
+    service_model = getattr(meta, "service_model", None)
+    service_name = getattr(service_model, "service_name", None)
+    return service_name == "bedrock-runtime"
+
+
+def _build_skip_dirs(client: Any) -> tuple[str, ...]:
+    """Add botocore / boto3 / aiobotocore to the stack-walker skip list so
+    attribution lands on the customer's code rather than on AWS SDK frames.
+    """
+    import frugal_metrics
+
+    candidates: list[Any] = [frugal_metrics]
+    for mod_name in ("botocore", "boto3", "aiobotocore"):
+        try:
+            mod = __import__(mod_name)
+        except ImportError:
+            continue
+        candidates.append(mod)
+    return tuple(
+        d for d in (caller_mod.package_dir_for(m) for m in candidates) if d
+    )
+
+
+def _build_call_context(
+    op_name: str,
+    kwargs: dict[str, Any],
+    skip_dirs: tuple[str, ...],
+) -> CallContext:
+    model_id = str(kwargs.get("modelId") or "<unknown>")
+    handler = bedrock_models.lookup_handler(model_id)
+    if op_name in _INVOKE_MODEL_OPS:
+        body_dict = bedrock_models.decode_invoke_model_body(kwargs.get("body"))
+        request_payload = handler.capture_invoke_model_payload(body_dict)
+    else:
+        request_payload = _capture_converse_payload(kwargs)
+    return CallContext(
+        system=semconv.SYSTEM_AWS_BEDROCK,
+        operation=semconv.OP_CHAT,
+        request_model=model_id,
+        skip_package_dirs=skip_dirs,
+        batched=False,
+        request_payload=request_payload,
+    )
+
+
+def _capture_converse_payload(kwargs: dict[str, Any]):
+    """Converse takes messages as top-level kwargs, no JSON body to decode."""
+    from frugal_metrics.prompt_analyzer import RequestPayload
+
+    inference_config = kwargs.get("inferenceConfig") or {}
+    max_tokens = inference_config.get("maxTokens")
+    if not isinstance(max_tokens, int):
+        max_tokens = None
+    tool_config = kwargs.get("toolConfig") or {}
+    return RequestPayload(
+        messages=kwargs.get("messages"),
+        tools=tool_config.get("tools"),
+        response_format=None,
+        max_tokens=max_tokens,
+        system=kwargs.get("system"),
+    )
+
+
+def _extract_non_streaming_usage(
+    op_name: str, kwargs: dict[str, Any], response: dict[str, Any]
+) -> UsageReport:
+    model_id = str(kwargs.get("modelId") or "")
+    handler = bedrock_models.lookup_handler(model_id)
+    usage = UsageReport()
+    if op_name == _OP_CONVERSE:
+        bedrock_models.merge_usage(
+            usage, handler.extract_converse_usage(response)
+        )
+        return usage
+    # invoke_model — try per-provider body parse first, then HTTP-header fallback
+    if handler is not bedrock_models.UNIVERSAL_FALLBACK_HANDLER:
+        body_dict = bedrock_models.consume_invoke_model_response_body(response)
+        bedrock_models.merge_usage(
+            usage, handler.extract_invoke_model_usage(body_dict)
+        )
+    bedrock_models.merge_usage(
+        usage, bedrock_models.extract_universal_headers_usage(response)
+    )
+    return usage
+
+
+def _make_stream_event_handler(op_name: str, model_id: str):
+    """Return the StreamEventHandler closure for the given operation + model."""
+    handler = bedrock_models.lookup_handler(model_id)
+
+    if op_name == _OP_CONVERSE_STREAM:
+        def handle_converse(event: Any, usage: UsageReport) -> None:
+            bedrock_models.merge_usage(
+                usage, bedrock_models.extract_converse_stream_event(event)
+            )
+        return handle_converse
+
+    # invoke_model_with_response_stream — per-provider extractor first,
+    # then universal final-chunk metrics so unrecognized models still get
+    # token counts.
+    def handle_invoke_model_stream(event: Any, usage: UsageReport) -> None:
+        bedrock_models.merge_usage(
+            usage, handler.extract_invoke_model_stream_event(event)
+        )
+        if handler is not bedrock_models.ANTHROPIC_CLAUDE_HANDLER:
+            # Anthropic's per-provider extractor already covers the
+            # message_start/message_delta path; the universal metrics block
+            # only appears in the terminal chunk so applying both for non-
+            # Anthropic models is safe and additive.
+            bedrock_models.merge_usage(
+                usage,
+                bedrock_models.extract_invoke_model_stream_event_universal(event),
+            )
+        else:
+            # For Anthropic specifically, only consult the universal block
+            # to fill in tokens that the per-provider extractor missed.
+            partial = bedrock_models.extract_invoke_model_stream_event_universal(event)
+            if partial is not None:
+                if usage.input_tokens is None:
+                    usage.input_tokens = partial.input_tokens
+                if usage.output_tokens is None:
+                    usage.output_tokens = partial.output_tokens
+
+    return handle_invoke_model_stream
+
+
+def _wrap_response_stream_sync(
+    op_name: str, kwargs: dict[str, Any], response: Any, emit: Any
+) -> Any:
+    model_id = str(kwargs.get("modelId") or "")
+    inner_stream = _resolve_stream_iterable(op_name, response)
+    handler_fn = _make_stream_event_handler(op_name, model_id)
+    proxy = SyncStreamProxy(inner_stream, emit, handler_fn)
+    response[_response_stream_key(op_name)] = proxy
+    return response
+
+
+def _wrap_response_stream_async(
+    op_name: str, kwargs: dict[str, Any], response: Any, emit: Any
+) -> Any:
+    model_id = str(kwargs.get("modelId") or "")
+    inner_stream = _resolve_stream_iterable(op_name, response)
+    handler_fn = _make_stream_event_handler(op_name, model_id)
+    proxy = AsyncStreamProxy(inner_stream, emit, handler_fn)
+    response[_response_stream_key(op_name)] = proxy
+    return response
+
+
+def _resolve_stream_iterable(op_name: str, response: dict[str, Any]) -> Any:
+    """Pick the right field on the boto3 response that holds the EventStream.
+
+    - converse_stream → response['stream']
+    - invoke_model_with_response_stream → response['body']
+    """
+    return response.get(_response_stream_key(op_name))
+
+
+def _response_stream_key(op_name: str) -> str:
+    return "stream" if op_name == _OP_CONVERSE_STREAM else "body"
+
+
+# --- The actual wrappers -----------------------------------------------------
+
+
+def _make_sync(
+    op_name: str,
+    original: Callable[..., Any],
+    skip_dirs: tuple[str, ...],
+    config: Any,
+) -> Callable[..., Any]:
+    @functools.wraps(original)
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
+        try:
+            ctx = _build_call_context(op_name, kwargs, skip_dirs)
+            emit = start_call(config, ctx)
+        except Exception as exc:  # noqa: BLE001
+            _record_internal_error("setup_failed", exc, config)
+            return original(*args, **kwargs)
+
+        try:
+            response = original(*args, **kwargs)
+        except BaseException as exc:
+            safe_call(lambda: emit.error(exc))
+            raise
+
+        if op_name in _STREAMING_OPS:
+            try:
+                return _wrap_response_stream_sync(op_name, kwargs, response, emit)
+            except Exception as exc:  # noqa: BLE001
+                _record_internal_error("stream_proxy_failed", exc, config)
+                return response
+
+        try:
+            usage = _extract_non_streaming_usage(op_name, kwargs, response)
+        except Exception:  # noqa: BLE001
+            usage = UsageReport()
+        safe_call(lambda: emit.success(usage))
+        return response
+
+    return wrapper
+
+
+def _make_async(
+    op_name: str,
+    original: Callable[..., Any],
+    skip_dirs: tuple[str, ...],
+    config: Any,
+) -> Callable[..., Any]:
+    @functools.wraps(original)
+    async def wrapper(*args: Any, **kwargs: Any) -> Any:
+        try:
+            ctx = _build_call_context(op_name, kwargs, skip_dirs)
+            emit = start_call(config, ctx)
+        except Exception as exc:  # noqa: BLE001
+            _record_internal_error("setup_failed", exc, config)
+            return await original(*args, **kwargs)
+
+        try:
+            response = await original(*args, **kwargs)
+        except BaseException as exc:
+            safe_call(lambda: emit.error(exc))
+            raise
+
+        if op_name in _STREAMING_OPS:
+            try:
+                return _wrap_response_stream_async(op_name, kwargs, response, emit)
+            except Exception as exc:  # noqa: BLE001
+                _record_internal_error("stream_proxy_failed", exc, config)
+                return response
+
+        try:
+            usage = _extract_non_streaming_usage(op_name, kwargs, response)
+        except Exception:  # noqa: BLE001
+            usage = UsageReport()
+        safe_call(lambda: emit.success(usage))
+        return response
+
+    return wrapper
+
+
+def _record_internal_error(stage: str, exc: BaseException, config: Any) -> None:
+    """Mirror of the same helper in wrappers/anthropic.py — emit the internal
+    error metric, swallow any failure trying to emit it.
+    """
+    try:
+        from frugal_metrics.otel import get_instruments
+
+        instruments = get_instruments(config)
+        err_attrs: dict[str, Any] = {
+            semconv.FRUGAL_PROJECT_ID: config.project_id,
+            "frugal.internal_error.stage": stage,
+            "frugal.internal_error.type": type(exc).__name__,
+        }
+        if config.customer_id:
+            err_attrs[semconv.FRUGAL_CUSTOMER_ID] = config.customer_id
+        instruments.internal_errors.add(1, err_attrs)
+    except Exception:  # noqa: BLE001
+        pass