Skip to content

Commit 65ec535

Browse files
chore(huggingface-hub): Fix chat completion stream type (#7070)
Add `huggingface-hub` to the typing dependency group and fix resulting mypy errors.
1 parent baacd47 commit 65ec535

3 files changed

Lines changed: 76 additions & 5 deletions

File tree

pyproject.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,7 @@ typing = [
7373
"pydantic>=2.13.4",
7474
"pydantic-ai-slim>=2.23.0",
7575
"langchain-core>=1.5.3",
76+
"huggingface-hub>=1.26.1",
7677
]
7778
test = [
7879
"dataclasses ; python_full_version < '3.7'",

sentry_sdk/integrations/huggingface_hub.py

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import inspect
22
import sys
33
from functools import wraps
4-
from typing import TYPE_CHECKING
4+
from typing import TYPE_CHECKING, cast
55

66
import sentry_sdk
77
from sentry_sdk.ai.monitoring import record_token_usage
@@ -24,6 +24,10 @@
2424
if TYPE_CHECKING:
2525
from typing import Any, Callable, Iterable, Union
2626

27+
from huggingface_hub import (
28+
ChatCompletionStreamOutput,
29+
)
30+
2731
from sentry_sdk.tracing import Span
2832

2933
try:
@@ -44,13 +48,13 @@ def __init__(
4448
@staticmethod
4549
def setup_once() -> None:
4650
# Other tasks that can be called: https://huggingface.co/docs/huggingface_hub/guides/inference#supported-providers-and-tasks
47-
huggingface_hub.inference._client.InferenceClient.text_generation = (
51+
huggingface_hub.inference._client.InferenceClient.text_generation = ( # type: ignore[method-assign]
4852
_wrap_huggingface_task(
4953
huggingface_hub.inference._client.InferenceClient.text_generation,
5054
OP.GEN_AI_TEXT_COMPLETION,
5155
)
5256
)
53-
huggingface_hub.inference._client.InferenceClient.chat_completion = (
57+
huggingface_hub.inference._client.InferenceClient.chat_completion = ( # type: ignore[method-assign]
5458
_wrap_huggingface_task(
5559
huggingface_hub.inference._client.InferenceClient.chat_completion,
5660
OP.GEN_AI_CHAT,
@@ -302,15 +306,15 @@ def new_details_iterator() -> "Iterable[Any]":
302306

303307
else:
304308
# chat-completion stream output
305-
def new_iterator() -> "Iterable[str]":
309+
def new_iterator() -> "Iterable[ChatCompletionStreamOutput]":
306310
finish_reason = None
307311
response_model = None
308312
response_text_buffer: "list[str]" = []
309313
tool_calls = None
310314
usage = None
311315

312316
with capture_internal_exceptions():
313-
for chunk in res:
317+
for chunk in cast("Iterable[ChatCompletionStreamOutput]", res):
314318
if hasattr(chunk, "model") and chunk.model is not None:
315319
response_model = chunk.model
316320

uv.lock

Lines changed: 66 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)