Coverage for haystack/components/routers/llm_messages_router.py: 99%
92 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-12 09:42 +0000
1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
2#
3# SPDX-License-Identifier: Apache-2.0
5import re
6from typing import Any
8from haystack import component, default_from_dict, default_to_dict
9from haystack.components.generators.chat.types import ChatGenerator
10from haystack.components.generators.utils import _trace_chat_generator_run
11from haystack.core.serialization import component_to_dict
12from haystack.dataclasses import ChatMessage, ChatRole
13from haystack.utils import deserialize_chatgenerator_inplace
14from haystack.utils.async_utils import _execute_component_async
16_RESERVED_OUTPUT_NAMES = ("chat_generator_text", "unmatched")
19@component
20class LLMMessagesRouter:
21 """
22 Routes Chat Messages to different connections using a generative Language Model to perform classification.
24 This component can be used with general-purpose LLMs and with specialized LLMs for moderation like Llama Guard.
26 ### Usage example
28 ```python
29 from haystack_integrations.components.generators.huggingface_api import HuggingFaceAPIChatGenerator
30 from haystack.components.routers.llm_messages_router import LLMMessagesRouter
31 from haystack.dataclasses import ChatMessage
33 # initialize a Chat Generator with a generative model for moderation
34 chat_generator = HuggingFaceAPIChatGenerator(
35 api_type="serverless_inference_api",
36 api_params={"model": "openai/gpt-oss-safeguard-20b", "provider": "groq"},
37 )
39 router = LLMMessagesRouter(chat_generator=chat_generator,
40 output_names=["unsafe", "safe"],
41 output_patterns=["unsafe", "safe"])
44 print(router.run([ChatMessage.from_user("How to rob a bank?")]))
46 # {
47 # 'chat_generator_text': 'unsafe\\nS2',
48 # 'unsafe': [
49 # ChatMessage(
50 # _role=<ChatRole.USER: 'user'>,
51 # _content=[TextContent(text='How to rob a bank?')],
52 # _name=None,
53 # _meta={}
54 # )
55 # ]
56 # }
57 ```
58 """
60 def __init__(
61 self,
62 chat_generator: ChatGenerator,
63 output_names: list[str],
64 output_patterns: list[str],
65 system_prompt: str | None = None,
66 ) -> None:
67 """
68 Initialize the LLMMessagesRouter component.
70 :param chat_generator: A ChatGenerator instance which represents the LLM.
71 :param output_names: A list of output connection names. These can be used to connect the router to other
72 components. `"chat_generator_text"` and `"unmatched"` are reserved output names and cannot be used here.
73 :param output_patterns: A list of regular expressions to be matched against the output of the LLM. Each pattern
74 corresponds to an output name. Patterns are evaluated in order.
75 When using moderation models, refer to the model card to understand the expected outputs.
76 :param system_prompt: An optional system prompt to customize the behavior of the LLM.
77 For moderation models, refer to the model card for supported customization options.
79 :raises ValueError: If output_names and output_patterns are not non-empty lists of the same length, or if
80 output_names uses the reserved names `"chat_generator_text"` or `"unmatched"`.
81 """
82 if not output_names or not output_patterns or len(output_names) != len(output_patterns):
83 raise ValueError("`output_names` and `output_patterns` must be non-empty lists of the same length")
85 collisions = set(_RESERVED_OUTPUT_NAMES).intersection(output_names)
86 if collisions:
87 names = ", ".join(repr(name) for name in sorted(collisions))
88 raise ValueError(
89 f"Output name(s) {names} are reserved output names in LLMMessagesRouter "
90 "('chat_generator_text' for the LLM reply, 'unmatched' for messages matching no pattern). "
91 "Rename the output to something else."
92 )
94 self._chat_generator = chat_generator
95 self._system_prompt = system_prompt
96 self._output_names = output_names
97 self._output_patterns = output_patterns
99 self._compiled_patterns = [re.compile(pattern) for pattern in output_patterns]
101 component.set_output_types(
102 self, **{"chat_generator_text": str, **dict.fromkeys(output_names + ["unmatched"], list[ChatMessage])}
103 )
105 def warm_up(self) -> None:
106 """Warm up the underlying chat generator."""
107 if hasattr(self._chat_generator, "warm_up"):
108 self._chat_generator.warm_up()
110 async def warm_up_async(self) -> None:
111 """Warm up the underlying chat generator on the serving event loop."""
112 if hasattr(self._chat_generator, "warm_up_async"):
113 await self._chat_generator.warm_up_async()
114 elif hasattr(self._chat_generator, "warm_up"):
115 self._chat_generator.warm_up()
117 def close(self) -> None:
118 """Release the underlying chat generator's resources."""
119 if hasattr(self._chat_generator, "close"):
120 self._chat_generator.close()
122 async def close_async(self) -> None:
123 """Release the underlying chat generator's async resources."""
124 if hasattr(self._chat_generator, "close_async"):
125 await self._chat_generator.close_async()
126 elif hasattr(self._chat_generator, "close"):
127 self._chat_generator.close()
129 def run(self, messages: list[ChatMessage]) -> dict[str, str | list[ChatMessage]]:
130 """
131 Classify the messages based on LLM output and route them to the appropriate output connection.
133 :param messages: A list of ChatMessages to be routed. Only user and assistant messages are supported.
135 :returns: A dictionary with the following keys:
136 - "chat_generator_text": The text output of the LLM, useful for debugging.
137 - "output_names": Each contains the list of messages that matched the corresponding pattern.
138 - "unmatched": The messages that did not match any of the output patterns.
140 :raises ValueError: If messages is an empty list or contains messages with unsupported roles.
141 """
142 if not messages:
143 raise ValueError("`messages` must be a non-empty list.")
144 if not all(message.is_from(ChatRole.USER) or message.is_from(ChatRole.ASSISTANT) for message in messages):
145 msg = (
146 "`messages` must contain only user and assistant messages. To customize the behavior of the "
147 "`chat_generator`, you can use the `system_prompt` parameter."
148 )
149 raise ValueError(msg)
151 self.warm_up()
153 messages_for_inference = []
154 if self._system_prompt:
155 messages_for_inference.append(ChatMessage.from_system(self._system_prompt))
156 messages_for_inference.extend(messages)
158 with _trace_chat_generator_run(self._chat_generator, {"messages": messages_for_inference}) as span:
159 generator_result = self._chat_generator.run(messages=messages_for_inference)
160 span.set_content_tag("haystack.component.output", generator_result)
161 chat_generator_text = generator_result["replies"][0].text
163 output = {"chat_generator_text": chat_generator_text}
165 for output_name, pattern in zip(self._output_names, self._compiled_patterns, strict=True):
166 if pattern.search(chat_generator_text):
167 output[output_name] = messages
168 break
169 else:
170 output["unmatched"] = messages
172 return output
174 async def run_async(self, messages: list[ChatMessage]) -> dict[str, str | list[ChatMessage]]:
175 """
176 Asynchronously classify the messages based on LLM output and route them to the appropriate output connection.
178 This is the asynchronous version of the `run` method. It has the same parameters and return values
179 but can be used with `await` in an async code. If the chat generator only implements a synchronous
180 `run` method, it is executed in a thread to avoid blocking the event loop.
182 :param messages: A list of ChatMessages to be routed. Only user and assistant messages are supported.
184 :returns: A dictionary with the following keys:
185 - "chat_generator_text": The text output of the LLM, useful for debugging.
186 - "output_names": Each contains the list of messages that matched the corresponding pattern.
187 - "unmatched": The messages that did not match any of the output patterns.
189 :raises ValueError: If messages is an empty list or contains messages with unsupported roles.
190 """
191 if not messages:
192 raise ValueError("`messages` must be a non-empty list.")
193 if not all(message.is_from(ChatRole.USER) or message.is_from(ChatRole.ASSISTANT) for message in messages):
194 msg = (
195 "`messages` must contain only user and assistant messages. To customize the behavior of the "
196 "`chat_generator`, you can use the `system_prompt` parameter."
197 )
198 raise ValueError(msg)
200 await self.warm_up_async()
202 messages_for_inference = []
203 if self._system_prompt:
204 messages_for_inference.append(ChatMessage.from_system(self._system_prompt))
205 messages_for_inference.extend(messages)
207 with _trace_chat_generator_run(self._chat_generator, {"messages": messages_for_inference}) as span:
208 generator_result = await _execute_component_async(self._chat_generator, messages=messages_for_inference)
209 span.set_content_tag("haystack.component.output", generator_result)
210 chat_generator_text = generator_result["replies"][0].text
212 output = {"chat_generator_text": chat_generator_text}
214 for output_name, pattern in zip(self._output_names, self._compiled_patterns, strict=True):
215 if pattern.search(chat_generator_text):
216 output[output_name] = messages
217 break
218 else:
219 output["unmatched"] = messages
221 return output
223 def to_dict(self) -> dict[str, Any]:
224 """
225 Serialize this component to a dictionary.
227 :returns:
228 The serialized component as a dictionary.
229 """
230 return default_to_dict(
231 self,
232 chat_generator=component_to_dict(obj=self._chat_generator, name="chat_generator"),
233 output_names=self._output_names,
234 output_patterns=self._output_patterns,
235 system_prompt=self._system_prompt,
236 )
238 @classmethod
239 def from_dict(cls, data: dict[str, Any]) -> "LLMMessagesRouter":
240 """
241 Deserialize this component from a dictionary.
243 :param data:
244 The dictionary representation of this component.
245 :returns:
246 The deserialized component instance.
247 """
248 if data["init_parameters"].get("chat_generator"):
249 deserialize_chatgenerator_inplace(data["init_parameters"], key="chat_generator")
251 return default_from_dict(cls, data)