Coverage for haystack/components/generators/chat/mock.py: 100%
154 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
5from __future__ import annotations
7import inspect
8import json
9import re
10from collections.abc import Callable, Sequence
11from dataclasses import replace
12from typing import Any
14from haystack import component, default_from_dict, default_to_dict, logging
15from haystack.components.generators.utils import _normalize_messages
16from haystack.dataclasses import (
17 ChatMessage,
18 ChatRole,
19 ComponentInfo,
20 FinishReason,
21 StreamingCallbackT,
22 StreamingChunk,
23 select_streaming_callback,
24)
25from haystack.dataclasses.streaming_chunk import ToolCallDelta, _invoke_streaming_callback
26from haystack.tools import ToolsType
27from haystack.utils import deserialize_callable, serialize_callable
29logger = logging.getLogger(__name__)
31# A callable that derives a response from the input messages, returning either the text of the assistant reply or a
32# full `ChatMessage`. Two shapes are supported and picked automatically from the callable's signature:
33# - `response_fn(messages)` — receives only the (normalized) list of input `ChatMessage` objects.
34# - `response_fn(messages, tools)` — additionally receives the `tools` passed to `run` (a `ToolsType` or `None`),
35# so the reply can depend on the runtime tool schema (for example, to build a tool call whose arguments follow
36# the tool's parameter names, or to route between the available tools).
37ResponseFn = Callable[..., str | ChatMessage]
40@component
41class MockChatGenerator:
42 """
43 A Chat Generator that returns predefined responses without calling any API.
45 It is a drop-in replacement for real Chat Generators (such as `OpenAIChatGenerator`) in tests, smoke tests, and
46 quick prototypes. It implements the same interface (`run`, `run_async`, streaming, serialization) but never
47 contacts an external service, so it is fully deterministic and free to run.
49 The response is selected based on how the component is configured:
51 - **Fixed response**: pass a single string or `ChatMessage`. The same reply is returned on every call.
52 Any `ChatMessage` passed as a response must have the `assistant` role.
53 - **Cycling responses**: pass a list of strings and/or `ChatMessage` objects. Each call returns the next item,
54 wrapping around to the start once the list is exhausted. This is useful to drive multi-step flows such as
55 Agents, where the first call returns a tool call and a later call returns the final answer.
56 - **Dynamic response**: pass a `response_fn` callable that receives the input messages and returns the reply.
57 This is useful when the reply should depend on the input, for example to echo back part of the prompt. If the
58 callable accepts a second positional argument, it also receives the `tools` passed to `run` (a `ToolsType` or
59 `None`), so the reply can depend on the runtime tool schema — handy for exercising Agents whose tool set varies.
60 - **Echo (default)**: with no configuration, the component echoes back the text of the last message that has
61 text content. This makes it usable out of the box for quick prototyping.
63 Pass `ChatMessage` objects (rather than plain strings) to return tool calls or reasoning content, which is handy
64 for exercising tool-calling pipelines without a real model.
66 ### Usage example
68 ```python
69 from haystack.components.generators.chat import MockChatGenerator
70 from haystack.dataclasses import ChatMessage, ToolCall
72 # Fixed response
73 generator = MockChatGenerator(responses="Hello, this is a mock response.")
74 result = generator.run([ChatMessage.from_user("Hi!")])
75 print(result["replies"][0].text) # "Hello, this is a mock response."
77 # Cycling responses to drive an Agent-like loop
78 generator = MockChatGenerator(
79 responses=[
80 ChatMessage.from_assistant(tool_calls=[ToolCall(tool_name="search", arguments={"query": "Haystack"})]),
81 "Here is the final answer.",
82 ]
83 )
85 # Dynamic, tool-aware response: build a tool call from the tools passed to run()
86 def call_first_tool(messages, tools):
87 if not tools:
88 return "No tools available."
89 return ChatMessage.from_assistant(
90 tool_calls=[ToolCall(tool_name=tools[0].name, arguments={})]
91 )
93 generator = MockChatGenerator(response_fn=call_first_tool)
94 ```
95 """
97 def __init__(
98 self,
99 responses: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
100 *,
101 response_fn: ResponseFn | None = None,
102 model: str = "mock-model",
103 meta: dict[str, Any] | None = None,
104 streaming_callback: StreamingCallbackT | None = None,
105 ) -> None:
106 """
107 Creates an instance of MockChatGenerator.
109 :param responses: The predefined response(s) to return. Accepts a single string or `ChatMessage` (returned on
110 every call), or a non-empty list of strings and/or `ChatMessage` objects that are returned in order,
111 cycling back to the start once exhausted. Strings are wrapped into assistant `ChatMessage` objects, and any
112 `ChatMessage` passed must have the `assistant` role. Mutually exclusive with `response_fn`. If neither is
113 provided, the component echoes the last message with text content.
114 :param response_fn: An optional callable that returns the reply as a string or an assistant `ChatMessage`. It
115 receives the input messages; if it accepts a second positional argument, it also receives the `tools`
116 passed to `run` (a `ToolsType` or `None`), letting the reply depend on the runtime tool schema. Use this
117 for input-dependent responses. Mutually exclusive with `responses`. To support serialization, pass a named
118 function (lambdas and nested functions cannot be serialized).
119 :param model: The model name reported in the response metadata. Purely cosmetic; no model is loaded.
120 :param meta: Additional metadata merged into the `meta` of every returned `ChatMessage`. A per-response
121 `ChatMessage`'s own metadata takes precedence over this value.
122 :param streaming_callback: An optional callback invoked with `StreamingChunk` objects reconstructed from the
123 predefined response. It lets the mock exercise streaming code paths without a real model.
124 :raises ValueError: If both `responses` and `response_fn` are provided, if `responses` is an empty list, or if
125 a `ChatMessage` response does not have the `assistant` role.
126 """
127 if responses is not None and response_fn is not None:
128 raise ValueError("Pass either 'responses' or 'response_fn', not both.")
130 self._responses = self._normalize_responses(responses)
131 self.response_fn = response_fn
132 self._response_fn_wants_tools = response_fn is not None and self._response_fn_accepts_tools(response_fn)
133 self.model = model
134 self.meta = meta or {}
135 self.streaming_callback = streaming_callback
136 self._call_count = 0
137 self._is_warmed_up = False
139 @staticmethod
140 def _normalize_responses(
141 responses: str | ChatMessage | Sequence[str | ChatMessage] | None,
142 ) -> list[ChatMessage] | None:
143 """Normalize the `responses` argument into a non-empty list of `ChatMessage`, or `None` for echo mode."""
144 if responses is None:
145 return None
147 items: list[str | ChatMessage]
148 if isinstance(responses, (str, ChatMessage)):
149 items = [responses]
150 elif isinstance(responses, Sequence):
151 items = list(responses)
152 else:
153 raise TypeError(f"'responses' must be a string, ChatMessage, or a sequence of them, got {type(responses)}.")
155 if len(items) == 0:
156 raise ValueError("'responses' must not be an empty list.")
158 normalized: list[ChatMessage] = []
159 for item in items:
160 if isinstance(item, str):
161 normalized.append(ChatMessage.from_assistant(item))
162 elif isinstance(item, ChatMessage):
163 if item.role != ChatRole.ASSISTANT:
164 raise ValueError(
165 f"Each ChatMessage response must have the 'assistant' role, got '{item.role.value}'."
166 )
167 normalized.append(item)
168 else:
169 raise TypeError(f"Each response must be a string or ChatMessage, got {type(item)}.")
170 return normalized
172 def to_dict(self) -> dict[str, Any]:
173 """Serialize the component to a dictionary."""
174 responses = [msg.to_dict() for msg in self._responses] if self._responses is not None else None
175 response_fn = serialize_callable(self.response_fn) if self.response_fn is not None else None
176 streaming_callback = serialize_callable(self.streaming_callback) if self.streaming_callback else None
177 return default_to_dict(
178 self,
179 responses=responses,
180 response_fn=response_fn,
181 model=self.model,
182 meta=self.meta,
183 streaming_callback=streaming_callback,
184 )
186 @classmethod
187 def from_dict(cls, data: dict[str, Any]) -> MockChatGenerator:
188 """Deserialize the component from a dictionary."""
189 init_params = data.get("init_parameters", {})
190 responses = init_params.get("responses")
191 if responses is not None:
192 init_params["responses"] = [ChatMessage.from_dict(msg) for msg in responses]
193 response_fn = init_params.get("response_fn")
194 if response_fn:
195 init_params["response_fn"] = deserialize_callable(response_fn)
196 streaming_callback = init_params.get("streaming_callback")
197 if streaming_callback:
198 init_params["streaming_callback"] = deserialize_callable(streaming_callback)
199 return default_from_dict(cls, data)
201 def warm_up(self) -> None:
202 """No-op warm up, provided for interface compatibility with real Chat Generators."""
203 self._is_warmed_up = True
205 @staticmethod
206 def _echo_text(messages: list[ChatMessage]) -> str | None:
207 """Return the text of the last message that has text content, for echo mode."""
208 for message in reversed(messages):
209 if message.text:
210 return message.text
211 return None
213 @staticmethod
214 def _response_fn_accepts_tools(response_fn: ResponseFn) -> bool:
215 """
216 Return True if `response_fn` can be called as `response_fn(messages, tools)`.
218 Callables that accept a single positional argument are called as `response_fn(messages)` instead.
219 """
220 try:
221 inspect.signature(response_fn).bind(None, None)
222 except (TypeError, ValueError):
223 # The callable rejects a second positional argument, or exposes no signature at all (some C callables).
224 return False
225 return True
227 @staticmethod
228 def _coerce_to_message(result: str | ChatMessage) -> ChatMessage:
229 """Turn the output of `response_fn` into a `ChatMessage`, wrapping strings and requiring the assistant role."""
230 if isinstance(result, str):
231 return ChatMessage.from_assistant(result)
232 if isinstance(result, ChatMessage):
233 if result.role != ChatRole.ASSISTANT:
234 raise ValueError(f"'response_fn' must return an assistant ChatMessage, got '{result.role.value}'.")
235 return result
236 raise TypeError(f"'response_fn' must return a string or ChatMessage, got {type(result)}.")
238 @staticmethod
239 def _estimate_usage(messages: list[ChatMessage], reply: ChatMessage) -> dict[str, int]:
240 """
241 Roughly estimate token usage as whitespace-separated word counts.
243 This is an approximation (not real tokenization) intended to give downstream code realistic-looking metadata.
244 """
245 prompt_tokens = sum(len((message.text or "").split()) for message in messages)
246 completion_tokens = len((reply.text or "").split())
247 return {
248 "prompt_tokens": prompt_tokens,
249 "completion_tokens": completion_tokens,
250 "total_tokens": prompt_tokens + completion_tokens,
251 }
253 def _build_meta(self, messages: list[ChatMessage], base: ChatMessage) -> dict[str, Any]:
254 """Build the metadata attached to the returned reply, merging defaults, init meta, and per-response meta."""
255 meta: dict[str, Any] = {
256 "model": self.model,
257 "index": 0,
258 "finish_reason": "tool_calls" if base.tool_calls else "stop",
259 "usage": self._estimate_usage(messages, base),
260 }
261 meta.update(self.meta)
262 meta.update(base.meta)
263 return meta
265 def _build_reply(self, messages: list[ChatMessage], tools: ToolsType | None) -> ChatMessage | None:
266 """Select and finalize the reply for the given input messages. Returns `None` when there is nothing to echo."""
267 if self.response_fn is not None:
268 raw = self.response_fn(messages, tools) if self._response_fn_wants_tools else self.response_fn(messages)
269 base = self._coerce_to_message(raw)
270 elif self._responses is not None:
271 base = self._responses[self._call_count % len(self._responses)]
272 self._call_count += 1
273 else:
274 text = self._echo_text(messages)
275 if text is None:
276 return None
277 base = ChatMessage.from_assistant(text)
279 return replace(base, _meta=self._build_meta(messages, base))
281 def _make_chunks(self, reply: ChatMessage) -> list[StreamingChunk]:
282 """Reconstruct streaming chunks from a finalized reply so streaming callbacks can be exercised."""
283 component_info = ComponentInfo.from_component(self)
284 chunks: list[StreamingChunk] = []
286 # Stream the text content word by word in content block 0.
287 parts = re.findall(r"\S+\s*", reply.text) if reply.text else []
288 for idx, part in enumerate(parts):
289 chunks.append(
290 StreamingChunk(
291 content=part, component_info=component_info, index=0, start=(idx == 0), meta={"model": self.model}
292 )
293 )
295 # Stream each tool call in its own content block.
296 block_index = 1 if parts else 0
297 for tool_call in reply.tool_calls:
298 chunks.append(
299 StreamingChunk(
300 content="",
301 component_info=component_info,
302 index=block_index,
303 start=True,
304 tool_calls=[
305 ToolCallDelta(
306 index=block_index,
307 tool_name=tool_call.tool_name,
308 arguments=json.dumps(tool_call.arguments),
309 id=tool_call.id,
310 )
311 ],
312 meta={"model": self.model},
313 )
314 )
315 block_index += 1
317 if not chunks:
318 chunks.append(
319 StreamingChunk(content="", component_info=component_info, index=0, meta={"model": self.model})
320 )
322 finish_reason: FinishReason = "tool_calls" if reply.tool_calls else "stop"
323 last = chunks[-1]
324 chunks[-1] = replace(last, finish_reason=finish_reason, meta={**last.meta, "finish_reason": finish_reason})
325 return chunks
327 @component.output_types(replies=list[ChatMessage])
328 def run(
329 self,
330 messages: list[ChatMessage] | str,
331 streaming_callback: StreamingCallbackT | None = None,
332 generation_kwargs: dict[str, Any] | None = None, # noqa: ARG002
333 *,
334 tools: ToolsType | None = None,
335 tools_strict: bool | None = None, # noqa: ARG002
336 ) -> dict[str, list[ChatMessage]]:
337 """
338 Return a predefined reply for the given messages without calling any API.
340 The signature mirrors `OpenAIChatGenerator.run` so the mock can be used as a positional drop-in replacement.
342 :param messages: The conversation history as a list of `ChatMessage` instances or a single string.
343 :param streaming_callback: An optional callback invoked with reconstructed `StreamingChunk` objects. Overrides
344 the callback set at initialization.
345 :param generation_kwargs: Accepted for interface compatibility and ignored.
346 :param tools: Passed to a tool-aware `response_fn` (one that accepts a second positional argument); otherwise
347 accepted for interface compatibility and ignored.
348 :param tools_strict: Accepted for interface compatibility and ignored.
349 :returns: A dictionary with a single key `replies` containing the predefined reply as a list of one
350 `ChatMessage` (empty in echo mode when there is no message to echo).
351 """
352 self.warm_up()
354 messages = _normalize_messages(messages)
355 streaming_callback = select_streaming_callback(
356 init_callback=self.streaming_callback, runtime_callback=streaming_callback, requires_async=False
357 )
359 reply = self._build_reply(messages, tools)
360 if reply is None:
361 return {"replies": []}
363 if streaming_callback is not None:
364 for chunk in self._make_chunks(reply):
365 streaming_callback(chunk)
367 return {"replies": [reply]}
369 @component.output_types(replies=list[ChatMessage])
370 async def run_async(
371 self,
372 messages: list[ChatMessage] | str,
373 streaming_callback: StreamingCallbackT | None = None,
374 generation_kwargs: dict[str, Any] | None = None, # noqa: ARG002
375 *,
376 tools: ToolsType | None = None,
377 tools_strict: bool | None = None, # noqa: ARG002
378 ) -> dict[str, list[ChatMessage]]:
379 """
380 Asynchronously return a predefined reply for the given messages without calling any API.
382 The signature mirrors `OpenAIChatGenerator.run_async` so the mock can be used as a positional drop-in
383 replacement.
385 :param messages: The conversation history as a list of `ChatMessage` instances or a single string.
386 :param streaming_callback: An optional callback invoked with reconstructed `StreamingChunk` objects. Overrides
387 the callback set at initialization.
388 :param generation_kwargs: Accepted for interface compatibility and ignored.
389 :param tools: Passed to a tool-aware `response_fn` (one that accepts a second positional argument); otherwise
390 accepted for interface compatibility and ignored.
391 :param tools_strict: Accepted for interface compatibility and ignored.
392 :returns: A dictionary with a single key `replies` containing the predefined reply as a list of one
393 `ChatMessage` (empty in echo mode when there is no message to echo).
394 """
395 if not self._is_warmed_up:
396 self.warm_up()
398 messages = _normalize_messages(messages)
399 streaming_callback = select_streaming_callback(
400 init_callback=self.streaming_callback, runtime_callback=streaming_callback, requires_async=True
401 )
403 reply = self._build_reply(messages, tools)
404 if reply is None:
405 return {"replies": []}
407 if streaming_callback is not None:
408 for chunk in self._make_chunks(reply):
409 await _invoke_streaming_callback(streaming_callback, chunk)
411 return {"replies": [reply]}