11import inspect
22import sys
33from functools import wraps
4- from typing import TYPE_CHECKING
4+ from typing import TYPE_CHECKING , cast
55
66import sentry_sdk
77from sentry_sdk .ai .monitoring import record_token_usage
2424if 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
2933try :
@@ -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
0 commit comments