from __future__ import annotations as _annotations from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator, Generator from contextlib import asynccontextmanager, contextmanager from dataclasses import dataclass, field from datetime import datetime, timezone from typing import Any, Literal, cast, overload from typing_extensions import assert_never from .. import ModelHTTPError, UnexpectedModelBehavior, _utils, usage from .._run_context import RunContext from .._thinking_part import split_content_into_text_and_thinking from .._utils import guard_tool_call_id as _guard_tool_call_id from ..messages import ( AudioUrl, BinaryContent, CachePoint, CompactionPart, DocumentUrl, FilePart, FinishReason, ImageUrl, ModelMessage, ModelRequest, ModelResponse, ModelResponsePart, ModelResponseStreamEvent, NativeToolCallPart, NativeToolReturnPart, RetryPromptPart, SpeechPart, SystemPromptPart, TextContent, TextPart, ThinkingPart, ToolAvailabilityDeltaPart, ToolCallPart, ToolReturnPart, UploadedFile, UserContent, UserPromptPart, VideoUrl, ) from ..profiles import DEFAULT_THINKING_TAGS, ModelProfile, ModelProfileSpec from ..providers import Provider, infer_provider from ..settings import ModelSettings from ..tools import ToolDefinition from . import ( Model, ModelRequestParameters, StreamedResponse, _unconverted_speech_part_error, # pyright: ignore[reportPrivateUsage] _unsynthesized_tool_availability_delta_error, # pyright: ignore[reportPrivateUsage] check_allow_model_requests, ) from ._tool_choice import resolve_tool_choice try: from huggingface_hub import ( AsyncInferenceClient, ChatCompletionInputFunctionName, ChatCompletionInputMessage, ChatCompletionInputMessageChunk, ChatCompletionInputTool, ChatCompletionInputToolCall, ChatCompletionInputToolChoiceClass, ChatCompletionInputURL, ChatCompletionOutput, ChatCompletionOutputMessage, ChatCompletionStreamOutput, TextGenerationOutputFinishReason, ) from huggingface_hub.errors import HfHubHTTPError except ImportError as _import_error: raise ImportError( 'Please install `huggingface_hub` to use Hugging Face Inference Providers, ' 'you can use the `huggingface` optional group — `pip install "pydantic-ai-slim[huggingface]"`' ) from _import_error @contextmanager def _map_api_errors(model_name: str) -> Generator[None]: try: yield except HfHubHTTPError as e: raise ModelHTTPError( status_code=e.response.status_code, model_name=model_name, body=e.response.content, headers=dict(e.response.headers), ) from e __all__ = ( 'HuggingFaceModel', 'HuggingFaceModelSettings', ) HFSystemPromptRole = Literal['system', 'user'] LatestHuggingFaceModelNames = Literal[ 'deepseek-ai/DeepSeek-R1', 'meta-llama/Llama-3.3-70B-Instruct', 'meta-llama/Llama-4-Maverick-17B-128E-Instruct', 'meta-llama/Llama-4-Scout-17B-16E-Instruct', 'Qwen/QwQ-32B', 'Qwen/Qwen2.5-72B-Instruct', 'Qwen/Qwen3-235B-A22B', 'Qwen/Qwen3-32B', ] """Latest Hugging Face models.""" HuggingFaceModelName = str | LatestHuggingFaceModelNames """Possible Hugging Face model names. You can browse available models [here](https://huggingface.co/models?pipeline_tag=text-generation&inference_provider=all&sort=trending). """ HuggingFaceFinishReason = Literal['stop', 'tool_calls'] | TextGenerationOutputFinishReason _FINISH_REASON_MAP: dict[HuggingFaceFinishReason, FinishReason] = { 'length': 'length', 'eos_token': 'stop', 'stop_sequence': 'stop', 'stop': 'stop', 'tool_calls': 'tool_call', } class HuggingFaceModelSettings(ModelSettings, total=False): """Settings used for a Hugging Face model request.""" # ALL FIELDS MUST BE `huggingface_` PREFIXED SO YOU CAN MERGE THEM WITH OTHER MODELS. # This class is a placeholder for any future huggingface-specific settings @dataclass(init=False) class HuggingFaceModel(Model[AsyncInferenceClient]): """A model that uses Hugging Face Inference Providers. Internally, this uses the [HF Python client](https://github.com/huggingface/huggingface_hub) to interact with the API. Apart from `__init__`, all methods are private or match those of the base class. """ _model_name: str = field(repr=False) _provider: Provider[AsyncInferenceClient] = field(repr=False) def __init__( self, model_name: str, *, provider: Literal['huggingface'] | Provider[AsyncInferenceClient] = 'huggingface', profile: ModelProfileSpec | None = None, settings: ModelSettings | None = None, ): """Initialize a Hugging Face model. Args: model_name: The name of the Model to use. You can browse available models [here](https://huggingface.co/models?pipeline_tag=text-generation&inference_provider=all&sort=trending). provider: The provider to use for Hugging Face Inference Providers. Can be either the string 'huggingface' or an instance of `Provider[AsyncInferenceClient]`. If not provided, the other parameters will be used. profile: The model profile to use. Defaults to a profile picked by the provider based on the model name. settings: Model-specific settings that will be used as defaults for this model. """ self._model_name = model_name if isinstance(provider, str): provider = infer_provider(provider) self._provider = provider super().__init__(settings=settings, profile=profile) @property def client(self) -> AsyncInferenceClient: return self._provider.client @property def base_url(self) -> str: """The base URL of the provider.""" return self._provider.base_url @property def model_name(self) -> HuggingFaceModelName: """The model name.""" return self._model_name @property def system(self) -> str: """The system / model provider.""" return self._provider.name async def request( self, messages: list[ModelMessage], model_settings: ModelSettings | None, model_request_parameters: ModelRequestParameters, ) -> ModelResponse: check_allow_model_requests() model_settings, model_request_parameters = self.prepare_request( model_settings, model_request_parameters, ) response = await self._completions_create( messages, False, cast(HuggingFaceModelSettings, model_settings or {}), model_request_parameters ) model_response = self._process_response(response) return model_response @asynccontextmanager async def request_stream( self, messages: list[ModelMessage], model_settings: ModelSettings | None, model_request_parameters: ModelRequestParameters, run_context: RunContext[Any] | None = None, ) -> AsyncGenerator[StreamedResponse]: check_allow_model_requests() model_settings, model_request_parameters = self.prepare_request( model_settings, model_request_parameters, ) response = await self._completions_create( messages, True, cast(HuggingFaceModelSettings, model_settings or {}), model_request_parameters ) try: yield await self._process_streamed_response(response, model_request_parameters) finally: aclose = getattr(response, 'aclose', None) if aclose is not None: # pragma: no branch await aclose() @overload async def _completions_create( self, messages: list[ModelMessage], stream: Literal[True], model_settings: HuggingFaceModelSettings, model_request_parameters: ModelRequestParameters, ) -> AsyncIterable[ChatCompletionStreamOutput]: ... @overload async def _completions_create( self, messages: list[ModelMessage], stream: Literal[False], model_settings: HuggingFaceModelSettings, model_request_parameters: ModelRequestParameters, ) -> ChatCompletionOutput: ... async def _completions_create( self, messages: list[ModelMessage], stream: bool, model_settings: HuggingFaceModelSettings, model_request_parameters: ModelRequestParameters, ) -> ChatCompletionOutput | AsyncIterable[ChatCompletionStreamOutput]: tools, tool_choice = self._get_tool_choice(model_settings, model_request_parameters) hf_messages = await self._map_messages(messages, model_request_parameters) with _map_api_errors(self.model_name): return await self.client.chat.completions.create( # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType, reportCallIssue] model=self._model_name, messages=hf_messages, # pyright: ignore[reportArgumentType] tools=tools, tool_choice=tool_choice or None, stream=stream, max_tokens=model_settings.get('max_tokens', None), stop=model_settings.get('stop_sequences', None), temperature=model_settings.get('temperature', None), top_p=model_settings.get('top_p', None), seed=model_settings.get('seed', None), presence_penalty=model_settings.get('presence_penalty', None), frequency_penalty=model_settings.get('frequency_penalty', None), logit_bias=model_settings.get('logit_bias', None), # pyright: ignore[reportArgumentType] logprobs=model_settings.get('logprobs', None), top_logprobs=model_settings.get('top_logprobs', None), extra_body=model_settings.get('extra_body'), # pyright: ignore[reportArgumentType] ) def _process_response(self, response: ChatCompletionOutput) -> ModelResponse: """Process a non-streamed response, and prepare a message to return.""" choice = response.choices[0] content = choice.message.content tool_calls = choice.message.tool_calls items: list[ModelResponsePart] = [] if content: items.extend( split_content_into_text_and_thinking(content, self.profile.get('thinking_tags', DEFAULT_THINKING_TAGS)) ) if tool_calls is not None: for c in tool_calls: items.append(ToolCallPart(c.function.name, c.function.arguments, tool_call_id=c.id)) raw_finish_reason = choice.finish_reason provider_details: dict[str, Any] = {'finish_reason': raw_finish_reason} if response.created: # pragma: no branch provider_details['timestamp'] = datetime.fromtimestamp(response.created, tz=timezone.utc) finish_reason = _FINISH_REASON_MAP.get(cast(HuggingFaceFinishReason, raw_finish_reason), None) return ModelResponse( parts=items, usage=_map_usage(response), model_name=response.model, provider_response_id=response.id, provider_name=self._provider.name, provider_url=self.base_url, finish_reason=finish_reason, provider_details=provider_details, ) async def _process_streamed_response( self, response: AsyncIterable[ChatCompletionStreamOutput], model_request_parameters: ModelRequestParameters ) -> StreamedResponse: """Process a streamed response, and prepare a streaming response to return.""" peekable_response: _utils.PeekableAsyncStream[ ChatCompletionStreamOutput, AsyncIterable[ChatCompletionStreamOutput] ] = _utils.PeekableAsyncStream(response) with _map_api_errors(self.model_name): first_chunk = await peekable_response.peek() if isinstance(first_chunk, _utils.Unset): raise UnexpectedModelBehavior( # pragma: no cover 'Streamed response ended without content or tool calls' ) # huggingface_hub types streaming responses as AsyncIterable, but the stream=True # response is an async generator at runtime. return HuggingFaceStreamedResponse( model_request_parameters=model_request_parameters, _model_name=first_chunk.model, _model_profile=self.profile, _response=peekable_response, _provider_name=self._provider.name, _provider_url=self.base_url, _provider_timestamp=datetime.fromtimestamp(first_chunk.created, tz=timezone.utc), ) @staticmethod def _get_tool_choice( model_settings: HuggingFaceModelSettings, model_request_parameters: ModelRequestParameters, ) -> tuple[ list[ChatCompletionInputTool], Literal['none', 'required', 'auto'] | ChatCompletionInputToolChoiceClass | None, ]: """Get tools and tool choice for the model. Returns a tuple of (tools, tool_choice). """ resolved_tool_choice = resolve_tool_choice(model_settings, model_request_parameters) tool_defs = model_request_parameters.declared_tool_defs tool_choice: Literal['none', 'required', 'auto'] | ChatCompletionInputToolChoiceClass | None if resolved_tool_choice in ('auto', 'required'): tool_choice = resolved_tool_choice elif resolved_tool_choice == 'none': # Use native 'none' mode to keep tool definitions cached while disabling tool calls tool_choice = 'none' elif isinstance(resolved_tool_choice, tuple): tool_choice_mode, tool_names = resolved_tool_choice if tool_choice_mode == 'required' and len(tool_names) == 1: tool_choice = ChatCompletionInputToolChoiceClass( function=ChatCompletionInputFunctionName(name=next(iter(tool_names))) ) else: # Breaks caching, but HuggingFace doesn't support limiting tools via API arg tool_defs = {k: v for k, v in tool_defs.items() if k in tool_names} tool_choice = tool_choice_mode else: assert_never(resolved_tool_choice) if not tool_defs: return [], None tools = [HuggingFaceModel._map_tool_definition(r) for r in tool_defs.values()] return tools, tool_choice async def _map_messages( # noqa: C901 self, messages: list[ModelMessage], model_request_parameters: ModelRequestParameters ) -> list[ChatCompletionInputMessage | ChatCompletionOutputMessage]: """Just maps a `pydantic_ai.Message` to a `huggingface_hub.ChatCompletionInputMessage`.""" hf_messages: list[ChatCompletionInputMessage | ChatCompletionOutputMessage] = [] for message in messages: if isinstance(message, ModelRequest): async for item in self._map_user_message(message): hf_messages.append(item) elif isinstance(message, ModelResponse): texts: list[str] = [] tool_calls: list[ChatCompletionInputToolCall] = [] for item in message.parts: if isinstance(item, TextPart): texts.append(item.content) elif isinstance(item, ToolCallPart): tool_calls.append(self._map_tool_call(item)) elif isinstance(item, ThinkingPart): start_tag, end_tag = self.profile.get('thinking_tags', DEFAULT_THINKING_TAGS) texts.append('\n'.join([start_tag, item.content, end_tag])) elif isinstance(item, NativeToolCallPart | NativeToolReturnPart): # pragma: no cover # This is currently never returned from huggingface pass elif isinstance(item, FilePart): # pragma: no cover # Files generated by models are not sent back to models that don't themselves generate files. pass elif isinstance(item, CompactionPart): # pragma: no cover # Compaction parts are not sent back to models that don't support compaction. pass elif isinstance(item, SpeechPart): # pragma: no cover # Unconverted realtime speech; `prepare_messages` turns these into `TextPart`s in `Model.prepare_messages`. raise _unconverted_speech_part_error() else: assert_never(item) message_param = ChatCompletionInputMessage(role='assistant') if texts: # Note: model responses from this model should only have one text item, so the following # shouldn't merge multiple texts into one unless you switch models between runs: message_param['content'] = '\n\n'.join(texts) if tool_calls: message_param['tool_calls'] = tool_calls hf_messages.append(message_param) else: assert_never(message) if instruction_parts := self._get_instruction_parts(messages, model_request_parameters): system_prompt_count = next( (i for i, m in enumerate(hf_messages) if getattr(m, 'role', None) != 'system'), len(hf_messages) ) hf_messages[system_prompt_count:system_prompt_count] = [ ChatCompletionInputMessage(content=part.content, role='system') for part in instruction_parts ] return hf_messages @staticmethod def _map_tool_call(t: ToolCallPart) -> ChatCompletionInputToolCall: return ChatCompletionInputToolCall.parse_obj_as_instance( # pyright: ignore[reportUnknownMemberType] { 'id': _guard_tool_call_id(t=t), 'type': 'function', 'function': { 'name': t.tool_name, 'arguments': t.args_as_json_str(), }, } ) @staticmethod def _map_tool_definition(f: ToolDefinition) -> ChatCompletionInputTool: tool_param: ChatCompletionInputTool = ChatCompletionInputTool.parse_obj_as_instance( # pyright: ignore[reportUnknownMemberType] { 'type': 'function', 'function': { 'name': f.name, 'description': f.description, 'parameters': f.parameters_json_schema, }, } ) return tool_param async def _map_user_message( self, message: ModelRequest ) -> AsyncIterable[ChatCompletionInputMessage | ChatCompletionOutputMessage]: file_content: list[UserContent] = [] for part in message.parts: if isinstance(part, SystemPromptPart): yield ChatCompletionInputMessage.parse_obj_as_instance({'role': 'system', 'content': part.content}) # pyright: ignore[reportUnknownMemberType] elif isinstance(part, UserPromptPart): yield await self._map_user_prompt(part) elif isinstance(part, ToolReturnPart): tool_text, tool_file_content = part.model_response_str_and_user_content() file_content.extend(tool_file_content) yield ChatCompletionOutputMessage.parse_obj_as_instance( # pyright: ignore[reportUnknownMemberType] { 'role': 'tool', 'tool_call_id': _guard_tool_call_id(t=part), 'content': tool_text, } ) elif isinstance(part, RetryPromptPart): if part.tool_name is None: yield ChatCompletionInputMessage.parse_obj_as_instance( # pyright: ignore[reportUnknownMemberType] {'role': 'user', 'content': part.model_response()} ) else: yield ChatCompletionInputMessage.parse_obj_as_instance( # pyright: ignore[reportUnknownMemberType] { 'role': 'tool', 'tool_call_id': _guard_tool_call_id(t=part), 'content': part.model_response(), } ) elif isinstance(part, ToolAvailabilityDeltaPart): # pragma: no cover raise _unsynthesized_tool_availability_delta_error() elif isinstance(part, SpeechPart): # pragma: no cover # Unconverted realtime speech; `prepare_messages` turns these into `UserPromptPart`s in `Model.prepare_messages`. raise _unconverted_speech_part_error() else: assert_never(part) if file_content: yield await self._map_user_prompt(UserPromptPart(content=file_content)) @staticmethod async def _map_user_prompt(part: UserPromptPart) -> ChatCompletionInputMessage: content: str | list[ChatCompletionInputMessageChunk] if isinstance(part.content, str): content = part.content else: content = [] for item in part.content: if isinstance(item, str | TextContent): text = item if isinstance(item, str) else item.content content.append(ChatCompletionInputMessageChunk(type='text', text=text)) elif isinstance(item, ImageUrl): url = ChatCompletionInputURL(url=item.url) content.append(ChatCompletionInputMessageChunk(type='image_url', image_url=url)) elif isinstance(item, BinaryContent): if item.is_image: url = ChatCompletionInputURL(url=item.data_uri) content.append(ChatCompletionInputMessageChunk(type='image_url', image_url=url)) else: # pragma: no cover raise RuntimeError(f'Unsupported binary content type: {item.media_type}') elif isinstance(item, AudioUrl): raise NotImplementedError('AudioUrl is not supported for Hugging Face') elif isinstance(item, DocumentUrl): raise NotImplementedError('DocumentUrl is not supported for Hugging Face') elif isinstance(item, VideoUrl): raise NotImplementedError('VideoUrl is not supported for Hugging Face') elif isinstance(item, UploadedFile): raise NotImplementedError('UploadedFile is not supported for Hugging Face') elif isinstance(item, CachePoint): # Hugging Face doesn't support prompt caching via CachePoint pass else: assert_never(item) return ChatCompletionInputMessage(role='user', content=content) @dataclass class HuggingFaceStreamedResponse(StreamedResponse): """Implementation of `StreamedResponse` for Hugging Face models.""" _model_name: str _model_profile: ModelProfile _response: _utils.PeekableAsyncStream[ChatCompletionStreamOutput, AsyncIterable[ChatCompletionStreamOutput]] _provider_name: str _provider_url: str _provider_timestamp: datetime | None = None _timestamp: datetime = field(default_factory=_utils.now_utc) async def close_stream(self) -> None: await self._response.aclose() async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]: with _map_api_errors(self._model_name): if self._provider_timestamp is not None: # pragma: no branch self.provider_details = {'timestamp': self._provider_timestamp} async for chunk in self._response: self._usage += _map_usage(chunk) if chunk.id: # pragma: no branch self.provider_response_id = chunk.id try: choice = chunk.choices[0] except IndexError: continue if raw_finish_reason := choice.finish_reason: self.provider_details = {**(self.provider_details or {}), 'finish_reason': raw_finish_reason} self.finish_reason = _FINISH_REASON_MAP.get(cast(HuggingFaceFinishReason, raw_finish_reason), None) # Handle the text part of the response content = choice.delta.content if content: for event in self._parts_manager.handle_text_delta( vendor_part_id='content', content=content, thinking_tags=self._model_profile.get('thinking_tags', DEFAULT_THINKING_TAGS), ignore_leading_whitespace=self._model_profile.get('ignore_streamed_leading_whitespace', False), ): yield event for dtc in choice.delta.tool_calls or []: # The SDK declares `function` as required, but the API can send chunks without it. maybe_event = self._parts_manager.handle_tool_call_delta( vendor_part_id=dtc.index, tool_name=dtc.function.name if dtc.function is not None else None, # pyright: ignore[reportUnnecessaryComparison] args=dtc.function.arguments if dtc.function is not None else None, # pyright: ignore[reportUnnecessaryComparison] tool_call_id=dtc.id, ) if maybe_event is not None: yield maybe_event @property def model_name(self) -> str: """Get the model name of the response.""" return self._model_name @property def provider_name(self) -> str: """Get the provider name.""" return self._provider_name @property def provider_url(self) -> str: """Get the provider base URL.""" return self._provider_url @property def timestamp(self) -> datetime: """Get the timestamp of the response.""" return self._timestamp def _map_usage(response: ChatCompletionOutput | ChatCompletionStreamOutput) -> usage.RequestUsage: response_usage = response.usage if response_usage is None: return usage.RequestUsage() return usage.RequestUsage( input_tokens=response_usage.prompt_tokens, output_tokens=response_usage.completion_tokens, )