Coverage for haystack/components/generators/utils.py: 90%
93 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-21 13:53 +0000
1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
2#
3# SPDX-License-Identifier: Apache-2.0
5import contextlib
6import json
7from collections.abc import Iterator
8from typing import Any
10from haystack import logging, tracing
11from haystack.dataclasses import ChatMessage, ReasoningContent, StreamingChunk, ToolCall
13logger = logging.getLogger(__name__)
16@contextlib.contextmanager
17def _trace_chat_generator_run(
18 chat_generator: Any, generator_inputs: dict[str, Any], parent_span: tracing.Span | None = None
19) -> Iterator[tracing.Span]:
20 """
21 Open a tracing span around a ChatGenerator call made internally by another component.
23 Components that embed a ChatGenerator but do not return its `ChatMessage` replies (for example rankers,
24 extractors or evaluators) would otherwise discard the LLM token usage carried in `reply.meta["usage"]`.
25 Wrapping the internal call in this span re-exposes that usage to tracers via the `haystack.component.output`
26 content tag, mirroring how the `Pipeline` traces its top-level components.
28 The caller is responsible for setting the output tag inside the context, so it is skipped when the call fails:
30 ```python
31 with _trace_chat_generator_run(self._chat_generator, {"messages": messages}) as span:
32 result = self._chat_generator.run(messages=messages)
33 span.set_content_tag("haystack.component.output", result)
34 ```
36 :param chat_generator: The ChatGenerator being invoked.
37 :param generator_inputs: The inputs passed to the generator, recorded as the span input content tag.
38 :param parent_span: Explicit parent span. Defaults to the current span. Pass it explicitly when the generator
39 runs in a worker thread, where the ambient span context does not propagate.
40 """
41 # Fall back to the active span so same-thread callers get correct nesting without passing a parent explicitly.
42 parent_span = parent_span or tracing.tracer.current_span()
43 with tracing.tracer.trace(
44 "haystack.chat_generator.run",
45 tags={"haystack.component.name": "chat_generator", "haystack.component.type": type(chat_generator).__name__},
46 parent_span=parent_span,
47 ) as span:
48 span.set_content_tag("haystack.component.input", generator_inputs)
49 yield span
52def print_streaming_chunk(chunk: StreamingChunk) -> None:
53 """
54 Callback function to handle and display streaming output chunks.
56 This function processes a `StreamingChunk` object by:
57 - Printing tool call metadata (if any), including function names and arguments, as they arrive.
58 - Printing tool call results when available.
59 - Printing the main content (e.g., text tokens) of the chunk as it is received.
61 The function outputs data directly to stdout and flushes output buffers to ensure immediate display during
62 streaming.
64 :param chunk: A chunk of streaming data containing content and optional metadata, such as tool calls and
65 tool results.
66 """
67 if chunk.start and chunk.index and chunk.index > 0:
68 # If this is the start of a new content block but not the first content block, print two new lines
69 print("\n\n", flush=True, end="")
71 ## Tool Call streaming
72 if chunk.tool_calls:
73 # Typically, if there are multiple tool calls in the chunk this means that the tool calls are fully formed and
74 # not just a delta.
75 for tool_call in chunk.tool_calls:
76 # If chunk.start is True indicates beginning of a tool call
77 # Also presence of tool_call.tool_name indicates the start of a tool call too
78 if chunk.start:
79 # If there is more than one tool call in the chunk, we print two new lines to separate them
80 # We know there is more than one tool call if the index of the tool call is greater than the index of
81 # the chunk.
82 if chunk.index and tool_call.index > chunk.index:
83 print("\n\n", flush=True, end="")
85 print(f"[TOOL CALL]\nTool: {tool_call.tool_name} \nArguments: ", flush=True, end="")
87 # print the tool arguments
88 if tool_call.arguments:
89 print(tool_call.arguments, flush=True, end="")
91 ## Tool Call Result streaming
92 # Print tool call results if available.
93 if chunk.tool_call_result:
94 # Tool Call Result is fully formed so delta accumulation is not needed
95 print(f"[TOOL RESULT]\n{chunk.tool_call_result.result}", flush=True, end="")
97 ## Normal content streaming
98 # Print the main content of the chunk (from ChatGenerator)
99 if chunk.content:
100 if chunk.start:
101 print("[ASSISTANT]\n", flush=True, end="")
102 print(chunk.content, flush=True, end="")
104 ## Reasoning content streaming
105 # Print the reasoning content of the chunk (from ChatGenerator)
106 if chunk.reasoning:
107 if chunk.start:
108 print("[REASONING]\n", flush=True, end="")
109 print(chunk.reasoning.reasoning_text, flush=True, end="")
111 # End of LLM assistant message so we add two new lines
112 # This ensures spacing between multiple LLM messages (e.g. Agent) or multiple Tool Call Results
113 if chunk.finish_reason is not None:
114 print("\n\n", flush=True, end="")
117def _convert_streaming_chunks_to_chat_message(chunks: list[StreamingChunk]) -> ChatMessage:
118 """
119 Connects the streaming chunks into a single ChatMessage.
121 :param chunks: The list of all `StreamingChunk` objects.
123 :returns: The ChatMessage.
124 """
125 text = "".join([chunk.content for chunk in chunks])
126 logprobs = []
127 for chunk in chunks:
128 if chunk.meta.get("logprobs"):
129 logprobs.append(chunk.meta.get("logprobs"))
130 tool_calls = []
132 # Accumulate reasoning content from chunks
133 reasoning_parts = [chunk.reasoning.reasoning_text for chunk in chunks if chunk.reasoning]
134 reasoning = ReasoningContent(reasoning_text="".join(reasoning_parts)) if reasoning_parts else None
136 # Process tool calls if present in any chunk
137 tool_call_data: dict[int, dict[str, str]] = {} # Track tool calls by index
138 for chunk in chunks:
139 if chunk.tool_calls:
140 for tool_call in chunk.tool_calls:
141 # We use the index of the tool_call to track the tool call across chunks since the ID is not always
142 # provided
143 if tool_call.index not in tool_call_data:
144 tool_call_data[tool_call.index] = {"id": "", "name": "", "arguments": ""}
146 # Save the ID if present
147 if tool_call.id is not None:
148 tool_call_data[tool_call.index]["id"] = tool_call.id
150 if tool_call.tool_name is not None:
151 tool_call_data[tool_call.index]["name"] += tool_call.tool_name
152 if tool_call.arguments is not None:
153 tool_call_data[tool_call.index]["arguments"] += tool_call.arguments
155 # Convert accumulated tool call data into ToolCall objects
156 sorted_keys = sorted(tool_call_data.keys())
157 for key in sorted_keys:
158 tool_call_dict = tool_call_data[key]
159 try:
160 arguments = json.loads(tool_call_dict.get("arguments", "{}")) if tool_call_dict.get("arguments") else {}
161 tool_calls.append(ToolCall(id=tool_call_dict["id"], tool_name=tool_call_dict["name"], arguments=arguments))
162 except json.JSONDecodeError:
163 logger.warning(
164 "The LLM provider returned a malformed JSON string for tool call arguments. This tool call "
165 "will be skipped. To always generate a valid JSON, set `tools_strict` to `True`. "
166 "Tool call ID: {_id}, Tool name: {_name}, Arguments: {_arguments}",
167 _id=tool_call_dict["id"],
168 _name=tool_call_dict["name"],
169 _arguments=tool_call_dict["arguments"],
170 )
172 # finish_reason can appear in different places so we look for the last one
173 finish_reasons = [chunk.finish_reason for chunk in chunks if chunk.finish_reason]
174 finish_reason = finish_reasons[-1] if finish_reasons else None
176 # usage info can appear in different chunks depending on the API provider
177 # (e.g., OpenAI returns it in the last chunk with empty choices, but Qwen3 may return it differently)
178 # so we look for the last non-None usage value across all chunks
179 usage = None
180 for chunk in reversed(chunks):
181 chunk_usage = chunk.meta.get("usage")
182 if chunk_usage is not None:
183 usage = chunk_usage
184 break
186 meta = {
187 "model": chunks[-1].meta.get("model"),
188 "index": 0,
189 "finish_reason": finish_reason,
190 "completion_start_time": chunks[0].meta.get("received_at"), # first chunk received
191 "usage": usage,
192 }
194 if logprobs:
195 meta["logprobs"] = logprobs
197 return ChatMessage.from_assistant(text=text or None, tool_calls=tool_calls, reasoning=reasoning, meta=meta)
200def _serialize_object(obj: Any) -> Any:
201 """
202 Convert an object to a serializable dict recursively.
204 Used to serialize `logprobs` and `usage` from OpenAI SDK response objects, so it skips any
205 attribute starting with "_" (SDK-internal fields). `base_serialization._serialize_value_with_schema`
206 doesn't skip those, so don't swap this out for it.
207 """
208 if hasattr(obj, "model_dump"):
209 return obj.model_dump()
210 if hasattr(obj, "__dict__"):
211 return {k: _serialize_object(v) for k, v in obj.__dict__.items() if not k.startswith("_")}
212 if isinstance(obj, dict):
213 return {k: _serialize_object(v) for k, v in obj.items()}
214 if isinstance(obj, list):
215 return [_serialize_object(item) for item in obj]
216 return obj
219def _normalize_messages(messages: list[ChatMessage] | str) -> list[ChatMessage]:
220 """Normalize messages to a list of ChatMessage objects."""
221 if isinstance(messages, str):
222 return [ChatMessage.from_user(messages)]
223 if isinstance(messages, list) and all(isinstance(msg, ChatMessage) for msg in messages):
224 return messages
225 raise TypeError("Invalid messages type. Expected list[ChatMessage] or str.")