fix(huggingface_hub): tolerate non-mapping responses

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

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),