From bcc585b442500cbdfd05e3c8f6be7914e2ca1a2d Mon Sep 17 00:00:00 2001 From: Starfolk Date: Fri, 17 Jul 2026 19:02:32 +0000 Subject: [PATCH] fix(huggingface_hub): drop eager _try_to_dict in stream tool-call merge, remove _get_field indirection MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - `_merge_tool_call_delta` no longer runs `_try_to_dict` (Pydantic `model_dump`) on every streaming tool-call delta chunk — `SpanImpl.log_internal` already runs `bt_safe_deep_copy` at send time. Same pattern as #586 (cohere) and #590 (google_genai). - Removed the `_get_field(obj, key)` dict-or-attr helper. HuggingFace's `BaseInferenceType` (parent of every response type and stream chunk on both matrix pins, `latest`=1.23.0 and floor `0.32.0`) inherits from `dict`, so `.get()` works uniformly on every response / chunk / delta. The two places that read from an `InferenceClient` instance now use `getattr(instance, key, None)` directly. - Net: 89 deletions / 60 insertions; also tightens the tool-call-delta loop to `isinstance(entry, dict)` and adds an explicit `if details is None: return {}` guard at the top of `_text_generation_extra_metadata` that was previously implicit via `_get_field`. Co-Authored-By: Claude Opus 4.7 --- .../integrations/huggingface_hub/tracing.py | 149 +++++++----------- 1 file changed, 60 insertions(+), 89 deletions(-) diff --git a/py/src/braintrust/integrations/huggingface_hub/tracing.py b/py/src/braintrust/integrations/huggingface_hub/tracing.py index 839d3345..29bcec81 100644 --- a/py/src/braintrust/integrations/huggingface_hub/tracing.py +++ b/py/src/braintrust/integrations/huggingface_hub/tracing.py @@ -32,7 +32,6 @@ _log_error_and_end_span, _normalize_chat_messages, _timing_metrics, - _try_to_dict, ) from braintrust.logger import start_span as _bt_start_span @@ -136,27 +135,14 @@ def start_span(*args, **kwargs): # --------------------------------------------------------------------------- # Field helpers +# +# HuggingFace inference response types (chat/text-generation/stream chunks) +# all inherit from ``dict`` via ``BaseInferenceType``, so ``.get()`` works +# uniformly on responses, chunks, and their nested fields. # --------------------------------------------------------------------------- -def _get_field(obj: Any, key: str) -> Any: - """Return a field from either a mapping or a HuggingFace SDK model object. - - The inference SDK returns dataclass-like objects (``BaseInferenceType``) - for most responses, but ``parse_obj_as_instance`` may surface them as - plain dicts when fields don't match the schema, and streaming chunks may - arrive as either. Centralizing the access keeps the rest of the module - indifferent to the runtime shape. - """ - if obj is None: - return None - if isinstance(obj, dict): - return obj.get(key) - return getattr(obj, key, None) - - def _first_nonempty_str(*candidates: Any, default: str | None = None) -> str | None: - """Return the first non-empty string in *candidates*, else *default*.""" for value in candidates: if isinstance(value, str) and value: return value @@ -164,23 +150,19 @@ def _first_nonempty_str(*candidates: Any, default: str | None = None) -> str | N def _resolve_provider_and_model(kwargs: dict[str, Any], instance: Any) -> tuple[str, str | None]: - """Return (provider, model) using request kwargs first, instance second. - - The HuggingFace Python SDK lets callers pin a default ``provider`` and - ``model`` on the ``InferenceClient`` instance and override them per-call. - Spans should mirror what actually went on the wire: per-call values win - over instance defaults, with ``"huggingface"`` as the final fallback for - provider so the integration identity is always present. ``model`` has no - such fallback — it stays ``None`` when neither side supplies one. + """Return (provider, model) with per-call kwargs winning over instance defaults. + + ``"huggingface"`` is the final fallback for provider so the integration + identity is always present; ``model`` has no such fallback. """ provider = _first_nonempty_str( kwargs.get("provider"), - _get_field(instance, "provider"), + getattr(instance, "provider", None), default=_PROVIDER, ) model = _first_nonempty_str( kwargs.get("model"), - _get_field(instance, "model"), + getattr(instance, "model", None), ) return provider, model @@ -218,13 +200,13 @@ def _extract_response_metadata(result: Any) -> dict[str, Any]: return {} metadata: dict[str, Any] = {} for key in _RESPONSE_METADATA_KEYS: - value = _get_field(result, key) + value = result.get(key) if value is not None: metadata[key] = value - choices = _get_field(result, "choices") + choices = result.get("choices") if isinstance(choices, list) and choices: - finish_reason = _get_field(choices[0], "finish_reason") + finish_reason = choices[0].get("finish_reason") if isinstance(finish_reason, str): metadata["finish_reason"] = finish_reason return metadata @@ -235,17 +217,17 @@ def _parse_usage_metrics(result: Any) -> dict[str, float]: if result is None: return {} - metrics: dict[str, float] = {} - usage = _get_field(result, "usage") + usage = result.get("usage") if usage is None: - return metrics + return {} + metrics: dict[str, float] = {} for key, metric in ( ("prompt_tokens", "prompt_tokens"), ("completion_tokens", "completion_tokens"), ("total_tokens", "tokens"), ): - value = _get_field(usage, key) + value = usage.get(key) if is_numeric(value): metrics[metric] = float(value) @@ -257,28 +239,23 @@ def _parse_usage_metrics(result: Any) -> dict[str, float]: def _text_generation_metrics(details: Any) -> dict[str, float]: """Extract metrics from a ``TextGenerationOutput.details`` payload. - - ``prompt_tokens`` from ``details.prefill`` length when available. - - ``completion_tokens`` from ``details.generated_tokens`` (or the length - of ``details.tokens`` if ``generated_tokens`` is missing). - - ``tokens`` = ``prompt_tokens + completion_tokens`` when either side is - known (missing side counted as 0). + ``prompt_tokens`` from ``details.prefill`` length, ``completion_tokens`` + from ``details.generated_tokens`` (falling back to ``len(details.tokens)`` + if the counter is missing), and ``tokens`` = their sum when either is + known. """ if details is None: return {} - prefill = _get_field(details, "prefill") - prompt_tokens: float | None = None - if isinstance(prefill, list): - prompt_tokens = float(len(prefill)) + prefill = details.get("prefill") + prompt_tokens: float | None = float(len(prefill)) if isinstance(prefill, list) else None - completion_tokens: float | None = None - generated_tokens = _get_field(details, "generated_tokens") + generated_tokens = details.get("generated_tokens") if is_numeric(generated_tokens): - completion_tokens = float(generated_tokens) + completion_tokens: float | None = float(generated_tokens) else: - tokens_list = _get_field(details, "tokens") - if isinstance(tokens_list, list): - completion_tokens = float(len(tokens_list)) + tokens_list = details.get("tokens") + completion_tokens = float(len(tokens_list)) if isinstance(tokens_list, list) else None metrics: dict[str, float] = {} if prompt_tokens is not None: @@ -298,32 +275,26 @@ def _text_generation_metrics(details: Any) -> dict[str, float]: def _chat_output(result: Any) -> Any: """Return the chat response's ``choices`` list verbatim. - The full ``choices`` array, with each entry's ``message`` / ``finish_reason`` - intact, is what the span logs as ``output``. This keeps tool calls, - logprobs, multiple choices, and any future fields available to consumers - without extra normalization. + Keeps tool calls, logprobs, multiple choices, and any future fields + available to consumers without extra normalization. """ if result is None: return None - choices = _get_field(result, "choices") - if not isinstance(choices, list): - return None - return choices + choices = result.get("choices") + return choices if isinstance(choices, list) else None def _text_generation_output(result: Any) -> Any: """Return ``{generated_text: ...}`` for any text-generation response shape. - Always emit an object so consumers can rely on a stable key regardless of - whether ``details=True`` was passed. When ``details=False`` the SDK - returns a plain ``str``; we wrap it. When ``details=True`` we pull - ``generated_text`` from the ``TextGenerationOutput``. + ``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} - generated_text = _get_field(result, "generated_text") + generated_text = result.get("generated_text") if isinstance(generated_text, str): return {"generated_text": generated_text} return result @@ -430,28 +401,25 @@ def _merge_tool_call_delta( """Merge a chat-completion streaming tool-call delta into the accumulator. Each delta contains ``index``, ``id``, ``type`` and a ``function`` block - whose ``arguments`` string is delivered piecewise across chunks. The - accumulator keeps function ``name`` once seen and concatenates the - incremental ``arguments`` strings. Tool-call entries are emitted in - ``sorted(index)`` order downstream, so no explicit order list is needed - here. + whose ``arguments`` string is delivered piecewise across chunks. The + accumulator concatenates the incremental ``arguments`` strings and keeps + function ``name`` once seen; entries are emitted in ``sorted(index)`` + order downstream. """ if not isinstance(delta_tool_calls, list): return for entry in delta_tool_calls: - entry_dict = entry if isinstance(entry, dict) else _try_to_dict(entry) - if not isinstance(entry_dict, dict): + if not isinstance(entry, dict): continue - index = entry_dict.get("index") + index = entry.get("index") if not isinstance(index, int): index = len(tool_calls_by_index) existing = tool_calls_by_index.get(index, {}) - # Carry forward static fields (id, type) when first seen. for key in ("id", "type"): - value = entry_dict.get(key) + value = entry.get(key) if value is not None: existing[key] = value - incoming_fn = entry_dict.get("function") + incoming_fn = entry.get("function") if isinstance(incoming_fn, dict): existing_fn = existing.get("function") or {} name = incoming_fn.get("name") @@ -459,7 +427,7 @@ def _merge_tool_call_delta( existing_fn["name"] = name args_delta = incoming_fn.get("arguments") if isinstance(args_delta, str): - existing_fn["arguments"] = f"{existing_fn.get('arguments', '') or ''}{args_delta}" + existing_fn["arguments"] = f"{existing_fn.get('arguments', '')}{args_delta}" existing["function"] = existing_fn tool_calls_by_index[index] = existing @@ -478,14 +446,14 @@ def record(self, chunk: Any) -> None: if chunk is None: return for key in _RESPONSE_METADATA_KEYS: - value = _get_field(chunk, key) + value = chunk.get(key) if value is not None and key not in self.metadata: self.metadata[key] = value - choices = _get_field(chunk, "choices") + choices = chunk.get("choices") if isinstance(choices, list): for choice in choices: - idx_raw = _get_field(choice, "index") + idx_raw = choice.get("index") index = idx_raw if isinstance(idx_raw, int) else 0 acc = self.choices_by_index.setdefault( index, @@ -497,18 +465,18 @@ def record(self, chunk: Any) -> None: }, ) - delta = _get_field(choice, "delta") - content = _get_field(delta, "content") + delta = choice.get("delta") or {} + content = delta.get("content") if isinstance(content, str) and content: acc["content_parts"].append(content) - role_value = _get_field(delta, "role") + role_value = delta.get("role") if isinstance(role_value, str): acc["role"] = role_value - _merge_tool_call_delta(acc["tool_calls_by_index"], _get_field(delta, "tool_calls")) + _merge_tool_call_delta(acc["tool_calls_by_index"], delta.get("tool_calls")) - fr = _get_field(choice, "finish_reason") + fr = choice.get("finish_reason") if isinstance(fr, str): acc["finish_reason"] = fr @@ -567,15 +535,16 @@ def record(self, chunk: Any) -> None: self.pieces.append(chunk) return - token_text = _get_field(_get_field(chunk, "token"), "text") + token = chunk.get("token") or {} + token_text = token.get("text") if isinstance(token_text, str): self.pieces.append(token_text) - generated_text = _get_field(chunk, "generated_text") + generated_text = chunk.get("generated_text") if isinstance(generated_text, str): self.final_text = generated_text - details = _get_field(chunk, "details") + details = chunk.get("details") if details is not None: self.last_details = details self.metadata.update(_text_generation_extra_metadata(details)) @@ -744,18 +713,20 @@ 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: + return {} metadata: dict[str, Any] = {} - finish_reason = _get_field(details, "finish_reason") + finish_reason = details.get("finish_reason") if isinstance(finish_reason, str): metadata["finish_reason"] = finish_reason - seed = _get_field(details, "seed") + seed = details.get("seed") if seed is not None: metadata["seed"] = seed return metadata def _log_text_generation_result(span, start_time: float, result: Any) -> None: - details = _get_field(result, "details") + details = result.get("details") if isinstance(result, dict) else None metrics = { **_timing_metrics(start_time, time.time()), **_text_generation_metrics(details),