diff --git a/py/src/braintrust/integrations/huggingface_hub/test_huggingface_hub.py b/py/src/braintrust/integrations/huggingface_hub/test_huggingface_hub.py index af57ff36b..01f6cda12 100644 --- a/py/src/braintrust/integrations/huggingface_hub/test_huggingface_hub.py +++ b/py/src/braintrust/integrations/huggingface_hub/test_huggingface_hub.py @@ -3,6 +3,7 @@ import asyncio import os import time +from collections import OrderedDict, UserDict import pytest from braintrust import logger, start_span @@ -496,7 +497,6 @@ def test_wrap_huggingface_hub_text_generation_details(memory_logger): @pytest.mark.vcr def test_wrap_huggingface_hub_feature_extraction_sync(memory_logger): pytest.importorskip("numpy") - assert not memory_logger.pop() client = wrap_huggingface_hub(_sync_client(model=EMBED_MODEL, provider=EMBED_PROVIDER)) @@ -670,3 +670,185 @@ async def _run(): class TestAutoInstrumentHuggingFaceHub: def test_auto_instrument_huggingface_hub(self): verify_autoinstrument_script("test_auto_huggingface_hub.py") + + +# --------------------------------------------------------------------------- +# VCR-backed integration tests (non-mapping response guard) +# +# These tests exercise the full client → wrapper → patcher → HTTP path +# with real cassettes. They verify that the instrumentation correctly +# handles real mapping responses end-to-end. Non-mapping edge cases +# (bytes, int, list) are covered by unit tests below — VCR cannot +# reproduce them because the real HF API always returns JSON mappings. +# --------------------------------------------------------------------------- + + +# --------------------------------------------------------------------------- +# Unit tests (non-mapping response guards) +# +# These tests call internal tracing functions directly with synthetic +# non-mapping inputs (bytes, int, list, None). They cannot use VCR because +# VCR records real HTTP traffic, and the real HF API always returns JSON +# mappings — there is no way to make it return b"raw video bytes" or 42 +# through a real HTTP call. These tests guard against the regression +# where a non-mapping response crashes the instrumentation. +# --------------------------------------------------------------------------- + + +def test_parse_usage_metrics_non_mapping_returns_no_metrics(): + """``_parse_usage_metrics`` is called from the chat and text-generation + logging paths, both of which will also be shared by the generative-media + wrappers (``text_to_video`` returns raw ``bytes``). A response that is not + mapping-like must degrade to no metrics rather than raising, so a successful + call is never turned into a traceback by its own instrumentation. + """ + from braintrust.integrations.huggingface_hub.tracing import ( + _parse_usage_metrics, + ) + + for value in [b"raw video bytes", b"", "a string", 42, ["a", "list"]]: + assert _parse_usage_metrics(value) == {} + + +def test_parse_usage_metrics_none_returns_no_metrics(): + from braintrust.integrations.huggingface_hub.tracing import ( + _parse_usage_metrics, + ) + + assert _parse_usage_metrics(None) == {} + + +def test_parse_usage_metrics_dict_without_usage_returns_no_metrics(): + from braintrust.integrations.huggingface_hub.tracing import ( + _parse_usage_metrics, + ) + + assert _parse_usage_metrics({"choices": []}) == {} + + +def test_parse_usage_metrics_dict_with_usage_is_unchanged(): + from braintrust.integrations.huggingface_hub.tracing import ( + _parse_usage_metrics, + ) + + assert _parse_usage_metrics({"usage": {"prompt_tokens": 3, "completion_tokens": 4}}) == { + "prompt_tokens": 3.0, + "completion_tokens": 4.0, + "tokens": 7.0, + } + + +def test_parse_usage_metrics_mapping_subclasses_still_yield_metrics(): + """Any ``Mapping`` must keep working, not just ``dict`` exactly. + + ``OrderedDict`` is a ``dict`` subclass while ``UserDict`` is only a + ``Mapping``, so a bare ``isinstance(result, dict)`` guard would accept + the former and silently drop token metrics for the latter. + """ + from braintrust.integrations.huggingface_hub.tracing import ( + _parse_usage_metrics, + ) + + for factory in [dict, OrderedDict, UserDict]: + payload = factory({"usage": {"prompt_tokens": 3, "completion_tokens": 4}}) + assert _parse_usage_metrics(payload) == { + "prompt_tokens": 3.0, + "completion_tokens": 4.0, + "tokens": 7.0, + } + + +def test_output_and_metadata_shapers_do_not_raise(): + """The chat and text-generation output/metadata shapers share the response + with ``_parse_usage_metrics``. Guarding only the metric parser relocates the + crash instead of removing it, so every shaper on that path is covered here. + """ + from braintrust.integrations.huggingface_hub.tracing import ( + _chat_output, + _extract_response_metadata, + _text_generation_extra_metadata, + _text_generation_output, + ) + + for value in [b"raw video bytes", 42, ["a", "list"], object()]: + assert _chat_output(value) is None + assert _extract_response_metadata(value) == {} + assert _text_generation_extra_metadata(value) == {} + assert _text_generation_output(value) is None + + +def test_text_generation_output_still_handles_str(): + from braintrust.integrations.huggingface_hub.tracing import ( + _text_generation_output, + ) + + assert _text_generation_output("plain text") == {"generated_text": "plain text"} + + +def test_log_chat_result_does_not_raise_on_bytes(memory_logger): + """Drive the full non-streaming chat logging path. + + ``_log_chat_result`` calls ``_parse_usage_metrics``, ``_chat_output`` + and ``_extract_response_metadata`` in sequence, so this fails if any one + of them is left unguarded. + """ + import time as _time + + from braintrust.integrations.huggingface_hub.tracing import _log_chat_result + + with start_span(name="huggingface.chat_completion") as span: + _log_chat_result(span, _time.time(), b"raw video bytes") + + # Reaching this point without an AttributeError is the assertion; the + # span is expected to be logged, with output/metadata simply empty. + spans = memory_logger.pop() + assert spans + + +def test_log_text_generation_result_accepts_mapping_subclass(memory_logger): + """Drive the full text-generation logging path with a ``Mapping``. + + ``_log_text_generation_result`` reads ``details`` behind an inline + ``isinstance(result, dict)`` guard. A ``UserDict`` is a ``Mapping`` but + not a ``dict``, so a ``dict`` guard silently drops the ``details`` + payload -- and with it the token metrics derived from it -- while every + other function on the path correctly accepts it. + """ + import time as _time + + from braintrust.integrations.huggingface_hub.tracing import ( + _log_text_generation_result, + ) + + payload = UserDict( + { + "generated_text": "hello", + "details": {"generated_tokens": 2}, + } + ) + + with start_span(name="huggingface.text_generation") as span: + _log_text_generation_result(span, _time.time(), payload) + + spans = memory_logger.pop() + assert spans + # The assertion that matters: token metrics must survive the path. With a + # ``dict`` guard the ``details`` payload is dropped and these are absent. + logged = spans[-1] + assert logged["metrics"].get("completion_tokens") == 2.0 + assert logged["metrics"].get("tokens") == 2.0 + + +def test_log_text_generation_result_does_not_raise_on_bytes(memory_logger): + """The same path must survive a non-mapping, non-``str`` response.""" + import time as _time + + from braintrust.integrations.huggingface_hub.tracing import ( + _log_text_generation_result, + ) + + with start_span(name="huggingface.text_generation") as span: + _log_text_generation_result(span, _time.time(), b"raw video bytes") + + spans = memory_logger.pop() + assert spans diff --git a/py/src/braintrust/integrations/huggingface_hub/tracing.py b/py/src/braintrust/integrations/huggingface_hub/tracing.py index 3cf24b58a..a960ca914 100644 --- a/py/src/braintrust/integrations/huggingface_hub/tracing.py +++ b/py/src/braintrust/integrations/huggingface_hub/tracing.py @@ -25,6 +25,7 @@ import logging import time +from collections.abc import Mapping from typing import Any, Protocol from braintrust.integrations.utils import ( @@ -197,7 +198,7 @@ def _build_request_metadata( def _extract_response_metadata(result: Any) -> dict[str, Any]: - if result is None: + if not isinstance(result, Mapping): return {} metadata: dict[str, Any] = {} for key in _RESPONSE_METADATA_KEYS: @@ -215,7 +216,7 @@ def _extract_response_metadata(result: Any) -> dict[str, Any]: def _parse_usage_metrics(result: Any) -> dict[str, float]: """Extract token usage from a chat or text-generation response.""" - if result is None: + if not isinstance(result, Mapping): return {} usage = result.get("usage") @@ -279,7 +280,7 @@ def _chat_output(result: Any) -> Any: Keeps tool calls, logprobs, multiple choices, and any future fields available to consumers without extra normalization. """ - if result is None: + if not isinstance(result, Mapping): return None choices = result.get("choices") return choices if isinstance(choices, list) else None @@ -291,10 +292,10 @@ def _text_generation_output(result: Any) -> Any: ``details=False`` returns a plain ``str``; ``details=True`` returns a ``TextGenerationOutput``. Wrap both into a stable-shape dict. """ - if result is None: - return None if isinstance(result, str): return {"generated_text": result} + if not isinstance(result, Mapping): + return None generated_text = result.get("generated_text") if isinstance(generated_text, str): return {"generated_text": generated_text} @@ -709,7 +710,7 @@ def _text_generation_extra_metadata(details: Any) -> dict[str, Any]: Shared by the non-streaming and streaming code paths so the two stay in sync when new ``details`` fields are added. """ - if details is None: + if not isinstance(details, Mapping): return {} metadata: dict[str, Any] = {} finish_reason = details.get("finish_reason") @@ -722,7 +723,7 @@ def _text_generation_extra_metadata(details: Any) -> dict[str, Any]: def _log_text_generation_result(span, start_time: float, result: Any) -> None: - details = result.get("details") if isinstance(result, dict) else None + details = result.get("details") if isinstance(result, Mapping) else None metrics = { **_timing_metrics(start_time, time.time()), **_text_generation_metrics(details),