Coverage for haystack/components/routers/llm_messages_router.py: 99%
87 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 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
17@component
18class LLMMessagesRouter:
19 """
20 Routes Chat Messages to different connections using a generative Language Model to perform classification.
22 This component can be used with general-purpose LLMs and with specialized LLMs for moderation like Llama Guard.
24 ### Usage example
26 ```python
27 from haystack_integrations.components.generators.huggingface_api import HuggingFaceAPIChatGenerator
28 from haystack.components.routers.llm_messages_router import LLMMessagesRouter
29 from haystack.dataclasses import ChatMessage
31 # initialize a Chat Generator with a generative model for moderation
32 chat_generator = HuggingFaceAPIChatGenerator(
33 api_type="serverless_inference_api",
34 api_params={"model": "openai/gpt-oss-safeguard-20b", "provider": "groq"},
35 )
37 router = LLMMessagesRouter(chat_generator=chat_generator,
38 output_names=["unsafe", "safe"],
39 output_patterns=["unsafe", "safe"])
42 print(router.run([ChatMessage.from_user("How to rob a bank?")]))
44 # {
45 # 'chat_generator_text': 'unsafe\\nS2',
46 # 'unsafe': [
47 # ChatMessage(
48 # _role=<ChatRole.USER: 'user'>,
49 # _content=[TextContent(text='How to rob a bank?')],
50 # _name=None,
51 # _meta={}
52 # )
53 # ]
54 # }
55 ```
56 """
58 def __init__(
59 self,
60 chat_generator: ChatGenerator,
61 output_names: list[str],
62 output_patterns: list[str],
63 system_prompt: str | None = None,
64 ) -> None:
65 """
66 Initialize the LLMMessagesRouter component.
68 :param chat_generator: A ChatGenerator instance which represents the LLM.
69 :param output_names: A list of output connection names. These can be used to connect the router to other
70 components.
71 :param output_patterns: A list of regular expressions to be matched against the output of the LLM. Each pattern
72 corresponds to an output name. Patterns are evaluated in order.
73 When using moderation models, refer to the model card to understand the expected outputs.
74 :param system_prompt: An optional system prompt to customize the behavior of the LLM.
75 For moderation models, refer to the model card for supported customization options.
77 :raises ValueError: If output_names and output_patterns are not non-empty lists of the same length.
78 """
79 if not output_names or not output_patterns or len(output_names) != len(output_patterns):
80 raise ValueError("`output_names` and `output_patterns` must be non-empty lists of the same length")
82 self._chat_generator = chat_generator
83 self._system_prompt = system_prompt
84 self._output_names = output_names
85 self._output_patterns = output_patterns
87 self._compiled_patterns = [re.compile(pattern) for pattern in output_patterns]
89 component.set_output_types(
90 self, **{"chat_generator_text": str, **dict.fromkeys(output_names + ["unmatched"], list[ChatMessage])}
91 )
93 def warm_up(self) -> None:
94 """Warm up the underlying chat generator."""
95 if hasattr(self._chat_generator, "warm_up"):
96 self._chat_generator.warm_up()
98 async def warm_up_async(self) -> None:
99 """Warm up the underlying chat generator on the serving event loop."""
100 if hasattr(self._chat_generator, "warm_up_async"):
101 await self._chat_generator.warm_up_async()
102 elif hasattr(self._chat_generator, "warm_up"):
103 self._chat_generator.warm_up()
105 def close(self) -> None:
106 """Release the underlying chat generator's resources."""
107 if hasattr(self._chat_generator, "close"):
108 self._chat_generator.close()
110 async def close_async(self) -> None:
111 """Release the underlying chat generator's async resources."""
112 if hasattr(self._chat_generator, "close_async"):
113 await self._chat_generator.close_async()
114 elif hasattr(self._chat_generator, "close"):
115 self._chat_generator.close()
117 def run(self, messages: list[ChatMessage]) -> dict[str, str | list[ChatMessage]]:
118 """
119 Classify the messages based on LLM output and route them to the appropriate output connection.
121 :param messages: A list of ChatMessages to be routed. Only user and assistant messages are supported.
123 :returns: A dictionary with the following keys:
124 - "chat_generator_text": The text output of the LLM, useful for debugging.
125 - "output_names": Each contains the list of messages that matched the corresponding pattern.
126 - "unmatched": The messages that did not match any of the output patterns.
128 :raises ValueError: If messages is an empty list or contains messages with unsupported roles.
129 """
130 if not messages:
131 raise ValueError("`messages` must be a non-empty list.")
132 if not all(message.is_from(ChatRole.USER) or message.is_from(ChatRole.ASSISTANT) for message in messages):
133 msg = (
134 "`messages` must contain only user and assistant messages. To customize the behavior of the "
135 "`chat_generator`, you can use the `system_prompt` parameter."
136 )
137 raise ValueError(msg)
139 self.warm_up()
141 messages_for_inference = []
142 if self._system_prompt:
143 messages_for_inference.append(ChatMessage.from_system(self._system_prompt))
144 messages_for_inference.extend(messages)
146 with _trace_chat_generator_run(self._chat_generator, {"messages": messages_for_inference}) as span:
147 generator_result = self._chat_generator.run(messages=messages_for_inference)
148 span.set_content_tag("haystack.component.output", generator_result)
149 chat_generator_text = generator_result["replies"][0].text
151 output = {"chat_generator_text": chat_generator_text}
153 for output_name, pattern in zip(self._output_names, self._compiled_patterns, strict=True):
154 if pattern.search(chat_generator_text):
155 output[output_name] = messages
156 break
157 else:
158 output["unmatched"] = messages
160 return output
162 async def run_async(self, messages: list[ChatMessage]) -> dict[str, str | list[ChatMessage]]:
163 """
164 Asynchronously classify the messages based on LLM output and route them to the appropriate output connection.
166 This is the asynchronous version of the `run` method. It has the same parameters and return values
167 but can be used with `await` in an async code. If the chat generator only implements a synchronous
168 `run` method, it is executed in a thread to avoid blocking the event loop.
170 :param messages: A list of ChatMessages to be routed. Only user and assistant messages are supported.
172 :returns: A dictionary with the following keys:
173 - "chat_generator_text": The text output of the LLM, useful for debugging.
174 - "output_names": Each contains the list of messages that matched the corresponding pattern.
175 - "unmatched": The messages that did not match any of the output patterns.
177 :raises ValueError: If messages is an empty list or contains messages with unsupported roles.
178 """
179 if not messages:
180 raise ValueError("`messages` must be a non-empty list.")
181 if not all(message.is_from(ChatRole.USER) or message.is_from(ChatRole.ASSISTANT) for message in messages):
182 msg = (
183 "`messages` must contain only user and assistant messages. To customize the behavior of the "
184 "`chat_generator`, you can use the `system_prompt` parameter."
185 )
186 raise ValueError(msg)
188 await self.warm_up_async()
190 messages_for_inference = []
191 if self._system_prompt:
192 messages_for_inference.append(ChatMessage.from_system(self._system_prompt))
193 messages_for_inference.extend(messages)
195 with _trace_chat_generator_run(self._chat_generator, {"messages": messages_for_inference}) as span:
196 generator_result = await _execute_component_async(self._chat_generator, messages=messages_for_inference)
197 span.set_content_tag("haystack.component.output", generator_result)
198 chat_generator_text = generator_result["replies"][0].text
200 output = {"chat_generator_text": chat_generator_text}
202 for output_name, pattern in zip(self._output_names, self._compiled_patterns, strict=True):
203 if pattern.search(chat_generator_text):
204 output[output_name] = messages
205 break
206 else:
207 output["unmatched"] = messages
209 return output
211 def to_dict(self) -> dict[str, Any]:
212 """
213 Serialize this component to a dictionary.
215 :returns:
216 The serialized component as a dictionary.
217 """
218 return default_to_dict(
219 self,
220 chat_generator=component_to_dict(obj=self._chat_generator, name="chat_generator"),
221 output_names=self._output_names,
222 output_patterns=self._output_patterns,
223 system_prompt=self._system_prompt,
224 )
226 @classmethod
227 def from_dict(cls, data: dict[str, Any]) -> "LLMMessagesRouter":
228 """
229 Deserialize this component from a dictionary.
231 :param data:
232 The dictionary representation of this component.
233 :returns:
234 The deserialized component instance.
235 """
236 if data["init_parameters"].get("chat_generator"):
237 deserialize_chatgenerator_inplace(data["init_parameters"], key="chat_generator")
239 return default_from_dict(cls, data)