Coverage for haystack/components/generators/chat/llm.py: 98%
40 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 typing import Any, Literal
7from haystack import component, logging
8from haystack.components.agents.agent import Agent
9from haystack.components.generators.chat.types import ChatGenerator
10from haystack.core.serialization import component_to_dict, default_from_dict, default_to_dict
11from haystack.dataclasses import ChatMessage, StreamingCallbackT
12from haystack.utils.callable_serialization import deserialize_callable, serialize_callable
13from haystack.utils.deserialization import deserialize_component_inplace
15logger = logging.getLogger(__name__)
18@component
19class LLM(Agent):
20 """
21 A text generation component powered by a large language model.
23 The LLM component is a simplified version of the Agent that focuses solely on text generation
24 without tool usage. It processes messages and returns a single response from the language model.
26 ### Usage examples
27 ```python
28 from haystack.components.generators.chat import LLM
29 from haystack.components.generators.chat import OpenAIChatGenerator
31 llm = LLM(
32 chat_generator=OpenAIChatGenerator(),
33 system_prompt="You are a helpful translation assistant.",
34 user_prompt="Summarize the following document: {{ document }}",
35 required_variables=["document"],
36 )
38 result = llm.run(document="The weather is lovely today and the sun is shining. ")
39 print(result["last_message"].text)
40 ```
41 """
43 def __init__(
44 self,
45 *,
46 chat_generator: ChatGenerator,
47 system_prompt: str | None = None,
48 user_prompt: str | None = None,
49 required_variables: list[str] | Literal["*"] = "*",
50 streaming_callback: StreamingCallbackT | None = None,
51 ) -> None:
52 """
53 Initialize the LLM component.
55 :param chat_generator: An instance of the chat generator that the LLM should use.
56 :param system_prompt: System prompt for the LLM. Can be a plain string template or a Jinja2 message template.
57 :param user_prompt: User prompt for the LLM. This prompt is appended to the messages provided at
58 runtime. Can be a plain string template or a Jinja2 message template. If it contains template variables
59 (e.g., `{{ variable_name }}`), they become inputs to the component. If omitted or if there are no
60 template variables, `messages` must be provided at runtime instead.
61 :param required_variables:
62 Variables that must be provided as input to `user_prompt` or `system_prompt`.
63 If a variable listed as required is not provided, an exception is raised.
64 If set to `"*"`, all variables found in the prompt are required. Defaults to `"*"`.
65 Only relevant when `user_prompt` or `system_prompt` contains template variables.
66 :param streaming_callback: A callback that will be invoked when a response is streamed from the LLM.
67 :raises ValueError: If user_prompt contains template variables but required_variables is an empty list.
68 """
69 super(LLM, self).__init__( # noqa: UP008
70 chat_generator=chat_generator,
71 system_prompt=system_prompt,
72 user_prompt=user_prompt,
73 required_variables=required_variables,
74 streaming_callback=streaming_callback,
75 )
76 if self._user_chat_prompt_builder is None or len(self._user_chat_prompt_builder.variables) == 0:
77 # This means user_prompt is empty or has no template variables.
78 # To ensure properly scheduling we then require messages to be passed at runtime.
79 component.set_input_type(self, "messages", list[ChatMessage])
80 else:
81 # user prompt was provided with variables
82 if isinstance(required_variables, list) and len(required_variables) == 0:
83 raise ValueError(
84 "required_variables must not be empty. Set it to '*' to require all variables, "
85 "or provide a non-empty list of variable names."
86 )
87 component.set_input_type(self, "messages", list[ChatMessage], None)
89 # The Agent base class declares `step_count` and `tool_call_counts` as outputs, but an LLM never has tools
90 # and always runs exactly one step — those values are uninformative, so drop them from the public surface.
91 # `token_usage` is still meaningful and stays exposed.
92 component.set_output_types(
93 self, messages=list[ChatMessage], last_message=ChatMessage, token_usage=dict[str, Any]
94 )
96 def to_dict(self) -> dict[str, Any]:
97 """
98 Serialize the LLM component to a dictionary.
100 :return: Dictionary with serialized data.
101 """
102 return default_to_dict(
103 self,
104 chat_generator=component_to_dict(obj=self.chat_generator, name="chat_generator"),
105 system_prompt=self.system_prompt,
106 user_prompt=self.user_prompt,
107 required_variables=self.required_variables,
108 streaming_callback=serialize_callable(self.streaming_callback) if self.streaming_callback else None,
109 )
111 @classmethod
112 def from_dict(cls, data: dict[str, Any]) -> "LLM":
113 """
114 Deserialize the LLM from a dictionary.
116 :param data: Dictionary to deserialize from.
117 :return: Deserialized LLM instance.
118 """
119 init_params = data.get("init_parameters", {})
121 deserialize_component_inplace(init_params, key="chat_generator")
123 if init_params.get("streaming_callback") is not None:
124 init_params["streaming_callback"] = deserialize_callable(init_params["streaming_callback"])
126 return default_from_dict(cls, data)
128 def run( # type: ignore[override] # `messages` is in **kwargs to allow dynamic required/optional status
129 self,
130 *,
131 streaming_callback: StreamingCallbackT | None = None,
132 generation_kwargs: dict[str, Any] | None = None,
133 **kwargs: Any,
134 ) -> dict[str, Any]:
135 """
136 Process messages and generate a response from the language model.
138 :param messages: Optional list of ChatMessage objects to prepend to the conversation. Whether this is
139 required or optional depends on the `user_prompt` configuration: if `user_prompt` has no template
140 variables, `messages` must be provided. Passed via `**kwargs`.
141 :param streaming_callback: A callback that will be invoked when a response is streamed from the LLM.
142 :param generation_kwargs: Additional keyword arguments for the chat generator. These are merged per key
143 with the `generation_kwargs` passed at the chat generator's initialization: keys provided here take
144 precedence, keys set only at initialization are kept.
145 :param kwargs: Additional keyword arguments. These are used to fill template variables in `user_prompt` or
146 `system_prompt` (the keys must match template variable names).
147 :returns:
148 A dictionary with the following keys:
149 - "messages": List of all messages exchanged during the LLM's run.
150 - "last_message": The last message exchanged during the LLM's run.
151 - "token_usage": Token usage from the LLM call (e.g. prompt_tokens, completion_tokens). Empty if the
152 chat generator did not return usage data.
153 """
154 # `messages` is intentionally omitted from the signature so the framework can treat it as required
155 # or optional depending on init configuration. See __init__ for details.
156 messages = kwargs.pop("messages", None)
157 result = super(LLM, self).run( # noqa: UP008
158 messages=messages or [],
159 streaming_callback=streaming_callback,
160 generation_kwargs=generation_kwargs,
161 **kwargs,
162 )
163 # Inherited Agent-internal bookkeeping that isn't useful at the LLM surface.
164 result.pop("step_count", None)
165 result.pop("tool_call_counts", None)
166 return result
168 async def run_async( # type: ignore[override] # `messages` is in **kwargs to allow dynamic required/optional status
169 self,
170 *,
171 streaming_callback: StreamingCallbackT | None = None,
172 generation_kwargs: dict[str, Any] | None = None,
173 **kwargs: Any,
174 ) -> dict[str, Any]:
175 """
176 Asynchronously process messages and generate a response from the language model.
178 :param messages: Optional list of ChatMessage objects to prepend to the conversation. Whether this is
179 required or optional depends on the `user_prompt` configuration: if `user_prompt` has no template
180 variables, `messages` must be provided. Passed via `**kwargs`.
181 :param streaming_callback: An asynchronous callback that will be invoked when a response is streamed
182 from the LLM.
183 :param generation_kwargs: Additional keyword arguments for the chat generator. These are merged per key
184 with the `generation_kwargs` passed at the chat generator's initialization: keys provided here take
185 precedence, keys set only at initialization are kept.
186 :param kwargs: Additional keyword arguments. These are used to fill template variables in `user_prompt` or
187 `system_prompt` (the keys must match template variable names).
188 :returns:
189 A dictionary with the following keys:
190 - "messages": List of all messages exchanged during the LLM's run.
191 - "last_message": The last message exchanged during the LLM's run.
192 - "token_usage": Token usage from the LLM call (e.g. prompt_tokens, completion_tokens). Empty if the
193 chat generator did not return usage data.
194 """
195 # `messages` is intentionally omitted from the signature so the framework can treat it as required
196 # or optional depending on init configuration. See __init__ for details.
197 messages = kwargs.pop("messages", None)
198 result = await super(LLM, self).run_async( # noqa: UP008
199 messages=messages or [],
200 streaming_callback=streaming_callback,
201 generation_kwargs=generation_kwargs,
202 **kwargs,
203 )
204 # Inherited Agent-internal bookkeeping that isn't useful at the LLM surface.
205 result.pop("step_count", None)
206 result.pop("tool_call_counts", None)
207 return result