12358207483

Committed 16 Dec 2024 05:40PM UTC coverage: 90.547% (+0.07%) from 90.48%

Build # 12358207483

Build Type

Pull #8640

github

Committed by

web-flow

Commit Message

Merge 2370c2f22 into a5b57f4b1

Pull Request Pull Request #8640: feat!: new `ChatMessage`

Run Details

8199 of 9055 relevant lines covered (90.55%)

0.91 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

97.95

haystack/dataclasses/chat_message.py

# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0

import warnings
from dataclasses import asdict, dataclass, field
from enum import Enum
from typing import Any, Dict, List, Optional, Sequence, Union


class ChatRole(str, Enum):
    """
    Enumeration representing the roles within a chat.
    """

    #: The user role. A message from the user contains only text.
    USER = "user"

    #: The system role. A message from the system contains only text.
    SYSTEM = "system"

    #: The assistant role. A message from the assistant can contain text and Tool calls. It can also store metadata.
    ASSISTANT = "assistant"

    #: The tool role. A message from a tool contains the result of a Tool invocation.
    TOOL = "tool"

    #: The function role. Deprecated in favor of `TOOL`.
    FUNCTION = "function"

    @staticmethod
    def from_str(string: str) -> "ChatRole":
        """
        Convert a string to a ChatRole enum.
        """
        enum_map = {e.value: e for e in ChatRole}
        role = enum_map.get(string)
        if role is None:
            msg = f"Unknown chat role '{string}'. Supported roles are: {list(enum_map.keys())}"
            raise ValueError(msg)
        return role


@dataclass
class ToolCall:
    """
    Represents a Tool call prepared by the model, usually contained in an assistant message.

    :param id: The ID of the Tool call.
    :param tool_name: The name of the Tool to call.
    :param arguments: The arguments to call the Tool with.
    """

    tool_name: str
    arguments: Dict[str, Any]
    id: Optional[str] = None  # noqa: A003


@dataclass
class ToolCallResult:
    """
    Represents the result of a Tool invocation.

    :param result: The result of the Tool invocation.
    :param origin: The Tool call that produced this result.
    :param error: Whether the Tool invocation resulted in an error.
    """

    result: str
    origin: ToolCall
    error: bool


@dataclass
class TextContent:
    """
    The textual content of a chat message.

    :param text: The text content of the message.
    """

    text: str


ChatMessageContentT = Union[TextContent, ToolCall, ToolCallResult]


@dataclass
class ChatMessage:
    """
    Represents a message in a LLM chat conversation.

    Use the `from_assistant`, `from_user`, `from_system`, and `from_tool` class methods to create a ChatMessage.
    """

    _role: ChatRole
    _content: Sequence[ChatMessageContentT]
    _meta: Dict[str, Any] = field(default_factory=dict, hash=False)

    def __new__(cls, *args, **kwargs):
        """
        This method is reimplemented to make the changes to the `ChatMessage` dataclass more visible.
        """

        general_msg = (
            "Use the `from_assistant`, `from_user`, `from_system`, and `from_tool` class methods to create a "
            "ChatMessage. Head over to the documentation for more information about the new API and how to migrate:"
            " https://docs.haystack.deepset.ai/docs/data-classes#chatmessage"
        )

        if "role" in kwargs or "content" in kwargs or "meta" in kwargs or "name" in kwargs:
            raise TypeError(
                "The `role`, `content`, `meta`, and `name` parameters of `ChatMessage` have been removed. "
                f"{general_msg}"
            )

        if len(args) > 1 and not isinstance(args[1], (TextContent, ToolCall, ToolCallResult)):
            raise TypeError(
                "The `content` parameter of `ChatMessage` must be a `ChatMessageContentT` instance. " f"{general_msg}"
            )

        return super(ChatMessage, cls).__new__(cls)

    def __post_init__(self):
        if self._role == ChatRole.FUNCTION:
            msg = "The `FUNCTION` role has been deprecated in favor of `TOOL` and will be removed in 2.10.0. "
            warnings.warn(msg, DeprecationWarning)

    def __getattribute__(self, name):
        """
        This method is reimplemented to make the `content` attribute removal more visible.
        """
        if name == "content":
            msg = (
                "The `content` attribute of `ChatMessage` has been removed. "
                "Use the `text` property to access the textual value. "
                "Head over to the documentation for more information: "
                "https://docs.haystack.deepset.ai/docs/data-classes#chatmessage"
            )
            raise AttributeError(msg)
        return object.__getattribute__(self, name)

    def __len__(self):
        return len(self._content)

    @property
    def role(self) -> ChatRole:
        """
        Returns the role of the entity sending the message.
        """
        return self._role

    @property
    def meta(self) -> Dict[str, Any]:
        """
        Returns the metadata associated with the message.
        """
        return self._meta

    @property
    def texts(self) -> List[str]:
        """
        Returns the list of all texts contained in the message.
        """
        return [content.text for content in self._content if isinstance(content, TextContent)]

    @property
    def text(self) -> Optional[str]:
        """
        Returns the first text contained in the message.
        """
        if texts := self.texts:
            return texts[0]
        return None

    @property
    def tool_calls(self) -> List[ToolCall]:
        """
        Returns the list of all Tool calls contained in the message.
        """
        return [content for content in self._content if isinstance(content, ToolCall)]

    @property
    def tool_call(self) -> Optional[ToolCall]:
        """
        Returns the first Tool call contained in the message.
        """
        if tool_calls := self.tool_calls:
            return tool_calls[0]
        return None

    @property
    def tool_call_results(self) -> List[ToolCallResult]:
        """
        Returns the list of all Tool call results contained in the message.
        """
        return [content for content in self._content if isinstance(content, ToolCallResult)]

    @property
    def tool_call_result(self) -> Optional[ToolCallResult]:
        """
        Returns the first Tool call result contained in the message.
        """
        if tool_call_results := self.tool_call_results:
            return tool_call_results[0]
        return None

    def is_from(self, role: Union[ChatRole, str]) -> bool:
        """
        Check if the message is from a specific role.

        :param role: The role to check against.
        :returns: True if the message is from the specified role, False otherwise.
        """
        if isinstance(role, str):
            role = ChatRole.from_str(role)
        return self._role == role

    @classmethod
    def from_user(cls, text: str, meta: Optional[Dict[str, Any]] = None) -> "ChatMessage":
        """
        Create a message from the user.

        :param text: The text content of the message.
        :param meta: Additional metadata associated with the message.
        :returns: A new ChatMessage instance.
        """
        return cls(_role=ChatRole.USER, _content=[TextContent(text=text)], _meta=meta or {})

    @classmethod
    def from_system(cls, text: str, meta: Optional[Dict[str, Any]] = None) -> "ChatMessage":
        """
        Create a message from the system.

        :param text: The text content of the message.
        :param meta: Additional metadata associated with the message.
        :returns: A new ChatMessage instance.
        """
        return cls(_role=ChatRole.SYSTEM, _content=[TextContent(text=text)], _meta=meta or {})

    @classmethod
    def from_assistant(
        cls,
        text: Optional[str] = None,
        meta: Optional[Dict[str, Any]] = None,
        tool_calls: Optional[List[ToolCall]] = None,
    ) -> "ChatMessage":
        """
        Create a message from the assistant.

        :param text: The text content of the message.
        :param meta: Additional metadata associated with the message.
        :param tool_calls: The Tool calls to include in the message.
        :returns: A new ChatMessage instance.
        """
        content: List[ChatMessageContentT] = []
        if text is not None:
            content.append(TextContent(text=text))
        if tool_calls:
            content.extend(tool_calls)

        return cls(_role=ChatRole.ASSISTANT, _content=content, _meta=meta or {})

    @classmethod
    def from_tool(
        cls, tool_result: str, origin: ToolCall, error: bool = False, meta: Optional[Dict[str, Any]] = None
    ) -> "ChatMessage":
        """
        Create a message from a Tool.

        :param tool_result: The result of the Tool invocation.
        :param origin: The Tool call that produced this result.
        :param error: Whether the Tool invocation resulted in an error.
        :param meta: Additional metadata associated with the message.
        :returns: A new ChatMessage instance.
        """
        return cls(
            _role=ChatRole.TOOL,
            _content=[ToolCallResult(result=tool_result, origin=origin, error=error)],
            _meta=meta or {},
        )

    @classmethod
    def from_function(cls, content: str, name: str) -> "ChatMessage":
        """
        Create a message from a function call. Deprecated in favor of `from_tool`.

        :param content: The text content of the message.
        :param name: The name of the function being called.
        :returns: A new ChatMessage instance.
        """
        msg = (
            "The `from_function` method is deprecated and will be removed in version 2.10.0. "
            "Its behavior has changed: it now attempts to convert legacy function messages to tool messages. "
            "This conversion is not guaranteed to succeed in all scenarios. "
            "Please migrate to `ChatMessage.from_tool` and carefully verify the results if you "
            "continue to use this method."
        )
        warnings.warn(msg)

        return cls.from_tool(content, ToolCall(id=None, tool_name=name, arguments={}), error=False)

    def to_dict(self) -> Dict[str, Any]:
        """
        Converts ChatMessage into a dictionary.

        :returns:
            Serialized version of the object.
        """
        serialized: Dict[str, Any] = {}
        serialized["_role"] = self._role.value
        serialized["_meta"] = self._meta

        content: List[Dict[str, Any]] = []
        for part in self._content:
            if isinstance(part, TextContent):
                content.append({"text": part.text})
            elif isinstance(part, ToolCall):
                content.append({"tool_call": asdict(part)})
            elif isinstance(part, ToolCallResult):
                content.append({"tool_call_result": asdict(part)})
            else:
                raise TypeError(f"Unsupported type in ChatMessage content: `{type(part).__name__}` for `{part}`.")

        serialized["_content"] = content
        return serialized

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "ChatMessage":
        """
        Creates a new ChatMessage object from a dictionary.

        :param data:
            The dictionary to build the ChatMessage object.
        :returns:
            The created object.
        """
        if "role" in data or "content" in data or "meta" in data or "name" in data:
            raise TypeError(
                "The `role`, `content`, `meta`, and `name` parameters of `ChatMessage` have been removed. "
                "Head over to the documentation for more information about the new API and how to migrate: "
                "https://docs.haystack.deepset.ai/docs/data-classes#chatmessage"
            )

        data["_role"] = ChatRole(data["_role"])

        content: List[ChatMessageContentT] = []

        for part in data["_content"]:
            if "text" in part:
                content.append(TextContent(text=part["text"]))
            elif "tool_call" in part:
                content.append(ToolCall(**part["tool_call"]))
            elif "tool_call_result" in part:
                result = part["tool_call_result"]["result"]
                origin = ToolCall(**part["tool_call_result"]["origin"])
                error = part["tool_call_result"]["error"]
                tcr = ToolCallResult(result=result, origin=origin, error=error)
                content.append(tcr)
            else:
                raise ValueError(f"Unsupported content in serialized ChatMessage: `{part}`")

        data["_content"] = content

        return cls(**data)

1	# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
2	#
3	# SPDX-License-Identifier: Apache-2.0
4
5	import warnings	1✔
6	from dataclasses import asdict, dataclass, field	1✔
7	from enum import Enum	1✔
8	from typing import Any, Dict, List, Optional, Sequence, Union	1✔
9
10
11	class ChatRole(str, Enum):	1✔
12	"""
13	Enumeration representing the roles within a chat.
14	"""
15
16	#: The user role. A message from the user contains only text.
17	USER = "user"	1✔
18
19	#: The system role. A message from the system contains only text.
20	SYSTEM = "system"	1✔
21
22	#: The assistant role. A message from the assistant can contain text and Tool calls. It can also store metadata.
23	ASSISTANT = "assistant"	1✔
24
25	#: The tool role. A message from a tool contains the result of a Tool invocation.
26	TOOL = "tool"	1✔
27
28	#: The function role. Deprecated in favor of `TOOL`.
29	FUNCTION = "function"	1✔
30
31	@staticmethod	1✔
32	def from_str(string: str) -> "ChatRole":	1✔
33	"""
34	Convert a string to a ChatRole enum.
35	"""
36	enum_map = {e.value: e for e in ChatRole}	1✔
37	role = enum_map.get(string)	1✔
38	if role is None:	1✔
39	msg = f"Unknown chat role '{string}'. Supported roles are: {list(enum_map.keys())}"	×
40	raise ValueError(msg)	×
41	return role	1✔
42
43
44	@dataclass	1✔
45	class ToolCall:	1✔
46	"""
47	Represents a Tool call prepared by the model, usually contained in an assistant message.
48
49	:param id: The ID of the Tool call.
50	:param tool_name: The name of the Tool to call.
51	:param arguments: The arguments to call the Tool with.
52	"""
53
54	tool_name: str	1✔
55	arguments: Dict[str, Any]	1✔
56	id: Optional[str] = None # noqa: A003	1✔
57
58
59	@dataclass	1✔
60	class ToolCallResult:	1✔
61	"""
62	Represents the result of a Tool invocation.
63
64	:param result: The result of the Tool invocation.
65	:param origin: The Tool call that produced this result.
66	:param error: Whether the Tool invocation resulted in an error.
67	"""
68
69	result: str	1✔
70	origin: ToolCall	1✔
71	error: bool	1✔
72
73
74	@dataclass	1✔
75	class TextContent:	1✔
76	"""
77	The textual content of a chat message.
78
79	:param text: The text content of the message.
80	"""
81
82	text: str	1✔
83
84
85	ChatMessageContentT = Union[TextContent, ToolCall, ToolCallResult]	1✔
86
87
88	@dataclass	1✔
89	class ChatMessage:	1✔
90	"""
91	Represents a message in a LLM chat conversation.
92
93	Use the `from_assistant`, `from_user`, `from_system`, and `from_tool` class methods to create a ChatMessage.
94	"""
95
96	_role: ChatRole	1✔
97	_content: Sequence[ChatMessageContentT]	1✔
98	_meta: Dict[str, Any] = field(default_factory=dict, hash=False)	1✔
99
100	def __new__(cls, args, *kwargs):	1✔
101	"""
102	This method is reimplemented to make the changes to the `ChatMessage` dataclass more visible.
103	"""
104
105	general_msg = (	1✔
106	"Use the `from_assistant`, `from_user`, `from_system`, and `from_tool` class methods to create a "
107	"ChatMessage. Head over to the documentation for more information about the new API and how to migrate:"
108	" https://docs.haystack.deepset.ai/docs/data-classes#chatmessage"
109	)
110
111	if "role" in kwargs or "content" in kwargs or "meta" in kwargs or "name" in kwargs:	1✔
112	raise TypeError(	1✔
113	"The `role`, `content`, `meta`, and `name` parameters of `ChatMessage` have been removed. "
114	f"{general_msg}"
115	)
116
117	if len(args) > 1 and not isinstance(args[1], (TextContent, ToolCall, ToolCallResult)):	1✔
118	raise TypeError(	1✔
119	"The `content` parameter of `ChatMessage` must be a `ChatMessageContentT` instance. " f"{general_msg}"
120	)
121
122	return super(ChatMessage, cls).__new__(cls)	1✔
123
124	def __post_init__(self):	1✔
125	if self._role == ChatRole.FUNCTION:	1✔
126	msg = "The `FUNCTION` role has been deprecated in favor of `TOOL` and will be removed in 2.10.0. "	1✔
127	warnings.warn(msg, DeprecationWarning)	1✔
128
129	def __getattribute__(self, name):	1✔
130	"""
131	This method is reimplemented to make the `content` attribute removal more visible.
132	"""
133	if name == "content":	1✔
134	msg = (	1✔
135	"The `content` attribute of `ChatMessage` has been removed. "
136	"Use the `text` property to access the textual value. "
137	"Head over to the documentation for more information: "
138	"https://docs.haystack.deepset.ai/docs/data-classes#chatmessage"
139	)
140	raise AttributeError(msg)	1✔
141	return object.__getattribute__(self, name)	1✔
142
143	def __len__(self):	1✔
144	return len(self._content)	1✔
145
146	@property	1✔
147	def role(self) -> ChatRole:	1✔
148	"""
149	Returns the role of the entity sending the message.
150	"""
151	return self._role	1✔
152
153	@property	1✔
154	def meta(self) -> Dict[str, Any]:	1✔
155	"""
156	Returns the metadata associated with the message.
157	"""
158	return self._meta	1✔
159
160	@property	1✔
161	def texts(self) -> List[str]:	1✔
162	"""
163	Returns the list of all texts contained in the message.
164	"""
165	return [content.text for content in self._content if isinstance(content, TextContent)]	1✔
166
167	@property	1✔
168	def text(self) -> Optional[str]:	1✔
169	"""
170	Returns the first text contained in the message.
171	"""
172	if texts := self.texts:	1✔
173	return texts[0]	1✔
174	return None	1✔
175
176	@property	1✔
177	def tool_calls(self) -> List[ToolCall]:	1✔
178	"""
179	Returns the list of all Tool calls contained in the message.
180	"""
181	return [content for content in self._content if isinstance(content, ToolCall)]	1✔
182
183	@property	1✔
184	def tool_call(self) -> Optional[ToolCall]:	1✔
185	"""
186	Returns the first Tool call contained in the message.
187	"""
188	if tool_calls := self.tool_calls:	1✔
189	return tool_calls[0]	1✔
190	return None	1✔
191
192	@property	1✔
193	def tool_call_results(self) -> List[ToolCallResult]:	1✔
194	"""
195	Returns the list of all Tool call results contained in the message.
196	"""
197	return [content for content in self._content if isinstance(content, ToolCallResult)]	1✔
198
199	@property	1✔
200	def tool_call_result(self) -> Optional[ToolCallResult]:	1✔
201	"""
202	Returns the first Tool call result contained in the message.
203	"""
204	if tool_call_results := self.tool_call_results:	1✔
205	return tool_call_results[0]	1✔
206	return None	1✔
207
208	def is_from(self, role: Union[ChatRole, str]) -> bool:	1✔
209	"""
210	Check if the message is from a specific role.
211
212	:param role: The role to check against.
213	:returns: True if the message is from the specified role, False otherwise.
214	"""
215	if isinstance(role, str):	1✔
216	role = ChatRole.from_str(role)	1✔
217	return self._role == role	1✔
218
219	@classmethod	1✔
220	def from_user(cls, text: str, meta: Optional[Dict[str, Any]] = None) -> "ChatMessage":	1✔
221	"""
222	Create a message from the user.
223
224	:param text: The text content of the message.
225	:param meta: Additional metadata associated with the message.
226	:returns: A new ChatMessage instance.
227	"""
228	return cls(_role=ChatRole.USER, _content=[TextContent(text=text)], _meta=meta or {})	1✔
229
230	@classmethod	1✔
231	def from_system(cls, text: str, meta: Optional[Dict[str, Any]] = None) -> "ChatMessage":	1✔
232	"""
233	Create a message from the system.
234
235	:param text: The text content of the message.
236	:param meta: Additional metadata associated with the message.
237	:returns: A new ChatMessage instance.
238	"""
239	return cls(_role=ChatRole.SYSTEM, _content=[TextContent(text=text)], _meta=meta or {})	1✔
240
241	@classmethod	1✔
242	def from_assistant(	1✔
243	cls,
244	text: Optional[str] = None,
245	meta: Optional[Dict[str, Any]] = None,
246	tool_calls: Optional[List[ToolCall]] = None,
247	) -> "ChatMessage":
248	"""
249	Create a message from the assistant.
250
251	:param text: The text content of the message.
252	:param meta: Additional metadata associated with the message.
253	:param tool_calls: The Tool calls to include in the message.
254	:returns: A new ChatMessage instance.
255	"""
256	content: List[ChatMessageContentT] = []	1✔
257	if text is not None:	1✔
258	content.append(TextContent(text=text))	1✔
259	if tool_calls:	1✔
260	content.extend(tool_calls)	1✔
261
262	return cls(_role=ChatRole.ASSISTANT, _content=content, _meta=meta or {})	1✔
263
264	@classmethod	1✔
265	def from_tool(	1✔
266	cls, tool_result: str, origin: ToolCall, error: bool = False, meta: Optional[Dict[str, Any]] = None
267	) -> "ChatMessage":
268	"""
269	Create a message from a Tool.
270
271	:param tool_result: The result of the Tool invocation.
272	:param origin: The Tool call that produced this result.
273	:param error: Whether the Tool invocation resulted in an error.
274	:param meta: Additional metadata associated with the message.
275	:returns: A new ChatMessage instance.
276	"""
277	return cls(	1✔
278	_role=ChatRole.TOOL,
279	_content=[ToolCallResult(result=tool_result, origin=origin, error=error)],
280	_meta=meta or {},
281	)
282
283	@classmethod	1✔
284	def from_function(cls, content: str, name: str) -> "ChatMessage":	1✔
285	"""
286	Create a message from a function call. Deprecated in favor of `from_tool`.
287
288	:param content: The text content of the message.
289	:param name: The name of the function being called.
290	:returns: A new ChatMessage instance.
291	"""
292	msg = (	1✔
293	"The `from_function` method is deprecated and will be removed in version 2.10.0. "
294	"Its behavior has changed: it now attempts to convert legacy function messages to tool messages. "
295	"This conversion is not guaranteed to succeed in all scenarios. "
296	"Please migrate to `ChatMessage.from_tool` and carefully verify the results if you "
297	"continue to use this method."
298	)
299	warnings.warn(msg)	1✔
300
301	return cls.from_tool(content, ToolCall(id=None, tool_name=name, arguments={}), error=False)	1✔
302
303	def to_dict(self) -> Dict[str, Any]:	1✔
304	"""
305	Converts ChatMessage into a dictionary.
306
307	:returns:
308	Serialized version of the object.
309	"""
310	serialized: Dict[str, Any] = {}	1✔
311	serialized["_role"] = self._role.value	1✔
312	serialized["_meta"] = self._meta	1✔
313
314	content: List[Dict[str, Any]] = []	1✔
315	for part in self._content:	1✔
316	if isinstance(part, TextContent):	1✔
317	content.append({"text": part.text})	1✔
318	elif isinstance(part, ToolCall):	1✔
319	content.append({"tool_call": asdict(part)})	1✔
320	elif isinstance(part, ToolCallResult):	1✔
321	content.append({"tool_call_result": asdict(part)})	1✔
322	else:
323	raise TypeError(f"Unsupported type in ChatMessage content: `{type(part).__name__}` for `{part}`.")	1✔
324
325	serialized["_content"] = content	1✔
326	return serialized	1✔
327
328	@classmethod	1✔
329	def from_dict(cls, data: Dict[str, Any]) -> "ChatMessage":	1✔
330	"""
331	Creates a new ChatMessage object from a dictionary.
332
333	:param data:
334	The dictionary to build the ChatMessage object.
335	:returns:
336	The created object.
337	"""
338	if "role" in data or "content" in data or "meta" in data or "name" in data:	1✔
339	raise TypeError(	×
340	"The `role`, `content`, `meta`, and `name` parameters of `ChatMessage` have been removed. "
341	"Head over to the documentation for more information about the new API and how to migrate: "
342	"https://docs.haystack.deepset.ai/docs/data-classes#chatmessage"
343	)
344
345	data["_role"] = ChatRole(data["_role"])	1✔
346
347	content: List[ChatMessageContentT] = []	1✔
348
349	for part in data["_content"]:	1✔
350	if "text" in part:	1✔
351	content.append(TextContent(text=part["text"]))	1✔
352	elif "tool_call" in part:	1✔
353	content.append(ToolCall(**part["tool_call"]))	1✔
354	elif "tool_call_result" in part:	1✔
355	result = part["tool_call_result"]["result"]	1✔
356	origin = ToolCall(**part["tool_call_result"]["origin"])	1✔
357	error = part["tool_call_result"]["error"]	1✔
358	tcr = ToolCallResult(result=result, origin=origin, error=error)	1✔
359	content.append(tcr)	1✔
360	else:
361	raise ValueError(f"Unsupported content in serialized ChatMessage: `{part}`")	1✔
362
363	data["_content"] = content	1✔
364
365	return cls(**data)	1✔

deepset-ai / haystack / 12358207483

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous