# # Copyright 2026 The Dapr Authors # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # http://www.apache.org/licenses/LICENSE-2.0 # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import dataclasses import logging from typing import Any, Callable, Dict, Iterator, Optional from huggingface_hub import ChatCompletionOutput, ChatCompletionStreamOutput from dapr_agents.types.message import ( AssistantMessage, FunctionCall, LLMChatCandidate, LLMChatCandidateChunk, LLMChatResponse, LLMChatResponseChunk, ToolCall, ToolCallChunk, ) logger = logging.getLogger(__name__) # Helper function to handle metadata extraction def _get_packet_metadata( pkt: Dict[str, Any], enrich_metadata: Optional[Dict[str, Any]] ) -> Dict[str, Any]: """ Extract metadata from HuggingFace packet and merge with enrich_metadata. Args: pkt (Dict[str, Any]): The HuggingFace packet from which to extract metadata. enrich_metadata (Optional[Dict[str, Any]]): Additional metadata to merge with the extracted metadata. Returns: Dict[str, Any]: The merged metadata dictionary. """ try: return { "id": pkt.get("id"), "created": pkt.get("created"), "model": pkt.get("model"), "object": pkt.get("object"), "service_tier": pkt.get("service_tier"), "system_fingerprint": pkt.get("system_fingerprint"), **(enrich_metadata or {}), } except Exception as e: logger.error(f"Failed to parse packet: {e}", exc_info=True) return {} # Helper function to process each choice delta (content, function call, tool call, finish reason) def _process_choice_delta( choice: Dict[str, Any], overall_meta: Dict[str, Any], on_chunk: Optional[Callable], first_chunk_flag: bool, ) -> Iterator[LLMChatResponseChunk]: """ Process each choice delta and yield corresponding chunks. Args: choice (Dict[str, Any]): The choice delta from HuggingFace response. overall_meta (Dict[str, Any]): Overall metadata to include in chunks. on_chunk (Optional[Callable]): Callback for each chunk. first_chunk_flag (bool): Flag indicating if this is the first chunk. Yields: LLMChatResponseChunk: The processed chunk with content, function call, tool calls, """ # Make an immutable snapshot for this single chunk meta = {**overall_meta} # mark first_chunk exactly once if first_chunk_flag and "first_chunk" not in meta: meta["first_chunk"] = True # Extract initial properties from choice delta: dict = choice.get("delta", {}) idx = choice.get("index") finish_reason = choice.get("finish_reason", None) logprobs = choice.get("logprobs", None) # Set additional metadata if finish_reason in ("stop", "tool_calls"): meta["last_chunk"] = True # Process content delta content = delta.get("content", None) function_call = delta.get("function_call", None) refusal = delta.get("refusal", None) role = delta.get("role", None) # Process tool calls chunk_tool_calls = [ToolCallChunk(**tc) for tc in (delta.get("tool_calls") or [])] # Initialize LLMChatResponseChunk response_chunk = LLMChatResponseChunk( result=LLMChatCandidateChunk( content=content, function_call=function_call, refusal=refusal, role=role, tool_calls=chunk_tool_calls, finish_reason=finish_reason, index=idx, logprobs=logprobs, ), metadata=meta, ) # Process chunk with on_chunk callback if on_chunk: on_chunk(response_chunk) # Yield LLMChatResponseChunk yield response_chunk # Main function to process HuggingFace streaming response def process_hf_stream( raw_stream: Iterator[ChatCompletionStreamOutput], *, enrich_metadata: Optional[Dict[str, Any]] = None, on_chunk: Optional[Callable], ) -> Iterator[LLMChatCandidateChunk]: """ Normalize HuggingFace streaming chat into LLMChatCandidateChunk objects, accumulating buffers per choice and yielding both partial and final chunks. Args: raw_stream: Iterator from client.chat.completions.create(..., stream=True) enrich_metadata: Extra key/value pairs to merge into each chunk.metadata on_chunk: Callback fired on every partial delta (token, function, tool) Yields: LLMChatCandidateChunk for every partial and final piece, in stream order """ enrich_metadata = enrich_metadata or {} overall_meta: Dict[str, Any] = {} # Track if we are in the first chunk first_chunk_flag = True for packet in raw_stream: # Convert Pydantic / HuggingFaceObject → plain dict if hasattr(packet, "model_dump"): pkt = packet.model_dump() elif hasattr(packet, "to_dict"): pkt = packet.to_dict() elif dataclasses.is_dataclass(packet): pkt = dataclasses.asdict(packet) else: raise TypeError(f"Cannot serialize packet of type {type(packet)}") # Capture overall metadata from the packet overall_meta = _get_packet_metadata(pkt, enrich_metadata) # Process each choice in this packet if choices := pkt.get("choices"): if len(choices) == 0: logger.warning( "Received empty 'choices' in HuggingFace packet, skipping." ) continue # Process the first choice in the packet choice = choices[0] yield from _process_choice_delta( choice, overall_meta, on_chunk, first_chunk_flag ) # Set first_chunk_flag to False after processing the first choice first_chunk_flag = False else: logger.warning(f" Yielding packet without 'choices': {pkt}") # Initialize default LLMChatResponseChunk final_response_chunk = LLMChatResponseChunk(metadata=overall_meta) yield final_response_chunk def process_hf_chat_response(response: ChatCompletionOutput) -> LLMChatResponse: """ Convert a non-streaming Hugging Face ChatCompletionOutput into our unified LLMChatResponse. This will: 1. Turn the HF dataclass into a plain dict via .model_dump() or .dict(). 2. Extract each `choice`, build an AssistantMessage (including any tool_calls or function_call shortcuts), wrap in LLMChatCandidate. 3. Collect top-level metadata (id, model, usage, etc.) into an LLMChatResponse. Args: response: The HFHubInferenceClientBase.chat.completions.create(...) output. Returns: An LLMChatResponse containing all chat candidates and metadata. """ # 1) serialise the HF object to a primitive dict try: if hasattr(response, "model_dump"): resp: Dict[str, Any] = response.model_dump() elif hasattr(response, "dict"): resp: Dict[str, Any] = response.dict() elif dataclasses.is_dataclass(response): resp = dataclasses.asdict(response) elif hasattr(response, "to_dict"): resp = response.to_dict() else: raise TypeError(f"Cannot serialize object of type {type(response)}") except Exception: logger.exception("Failed to serialize HF chat response") resp = {} candidates = [] for choice in resp.get("choices", []): msg = choice.get("message") or {} # 2) build tool_calls list if present tool_calls: Optional[list[ToolCall]] = None if msg.get("tool_calls"): tool_calls = [] for tc in msg["tool_calls"]: try: tool_calls.append(ToolCall(**tc)) except Exception: logger.exception(f"Invalid HF tool_call entry: {tc}") # 2b) handle the single‑ID shortcut if msg.get("tool_call_id") and not tool_calls: # HF only sent you an ID; we turn that into a zero‑arg function_call fc = FunctionCall(name=msg["tool_call_id"], arguments="") tool_calls = [ ToolCall(id=msg["tool_call_id"], type="function", function=fc) ] # 3) promote first tool_call into function_call if desired function_call = tool_calls[0].function if tool_calls else None assistant = AssistantMessage( content=msg.get("content"), refusal=None, tool_calls=tool_calls, function_call=function_call, ) candidates.append( LLMChatCandidate( message=assistant, finish_reason=choice.get("finish_reason"), index=choice.get("index"), logprobs=choice.get("logprobs"), ) ) # 4) collect overall metadata metadata = { "provider": "huggingface", "id": resp.get("id"), "model": resp.get("model"), "created": resp.get("created"), "system_fingerprint": resp.get("system_fingerprint"), "usage": resp.get("usage"), } return LLMChatResponse(results=candidates, metadata=metadata)