From 51c13904260604644edc2e9edc9969b54a06edc9 Mon Sep 17 00:00:00 2001 From: Stefano Fiorucci Date: Thu, 28 Nov 2024 10:35:24 +0100 Subject: [PATCH] chore: use class methods to create `ChatMessage` (#8581) * use class methods to build messages * fix failing format --- haystack/core/pipeline/base.py | 3 +- ...ge-use-class-methods-85a30b0c5512b08c.yaml | 5 + .../builders/test_answer_builder.py | 132 ++++++------------ .../builders/test_chat_prompt_builder.py | 9 +- test/dataclasses/test_chat_message.py | 8 +- 5 files changed, 62 insertions(+), 95 deletions(-) create mode 100644 releasenotes/notes/chatmessage-use-class-methods-85a30b0c5512b08c.yaml diff --git a/haystack/core/pipeline/base.py b/haystack/core/pipeline/base.py index 8304d5ff85..1db8230a3b 100644 --- a/haystack/core/pipeline/base.py +++ b/haystack/core/pipeline/base.py @@ -183,7 +183,8 @@ def from_dict( "Component was deserialized with the following data:" f"{str(component_data)}." "Possible reasons include malformed serialized data, mismatch between the serialized component " - "and the loaded one (due to a breaking change, see https://github.com/deepset-ai/haystack/releases), etc." + "and the loaded one (due to a breaking change, see " + "https://github.com/deepset-ai/haystack/releases), etc." ) raise DeserializationError(msg) from error pipe.add_component(name=name, instance=instance) diff --git a/releasenotes/notes/chatmessage-use-class-methods-85a30b0c5512b08c.yaml b/releasenotes/notes/chatmessage-use-class-methods-85a30b0c5512b08c.yaml new file mode 100644 index 0000000000..edc24f8ab4 --- /dev/null +++ b/releasenotes/notes/chatmessage-use-class-methods-85a30b0c5512b08c.yaml @@ -0,0 +1,5 @@ +--- +enhancements: + - | + Across Haystack codebase, we have replaced the use of `ChatMessage` dataclass constructor with specific + class methods (`ChatMessage.from_user`, `ChatMessage.from_assistant`, etc.). diff --git a/test/components/builders/test_answer_builder.py b/test/components/builders/test_answer_builder.py index 9c220ef477..73f20b6d97 100644 --- a/test/components/builders/test_answer_builder.py +++ b/test/components/builders/test_answer_builder.py @@ -164,77 +164,54 @@ def test_run_with_reference_pattern_set_at_runtime(self): def test_run_with_chat_message_replies_without_pattern(self): component = AnswerBuilder() - replies = [ - ChatMessage( - content="Answer: AnswerString", - role=ChatRole.ASSISTANT, - name=None, - meta={ - "model": "gpt-4o-mini", - "index": 0, - "finish_reason": "stop", - "usage": {"prompt_tokens": 32, "completion_tokens": 153, "total_tokens": 185}, - }, - ) - ] - output = component.run(query="test query", replies=replies, meta=[{}]) - answers = output["answers"] - assert len(answers) == 1 - assert answers[0].data == "Answer: AnswerString" - assert answers[0].meta == { + + message_meta = { "model": "gpt-4o-mini", "index": 0, "finish_reason": "stop", "usage": {"prompt_tokens": 32, "completion_tokens": 153, "total_tokens": 185}, } + replies = [ChatMessage.from_assistant("Answer: AnswerString", meta=message_meta)] + + output = component.run(query="test query", replies=replies, meta=[{}]) + answers = output["answers"] + assert len(answers) == 1 + assert answers[0].data == "Answer: AnswerString" + assert answers[0].meta == message_meta assert answers[0].query == "test query" assert answers[0].documents == [] assert isinstance(answers[0], GeneratedAnswer) def test_run_with_chat_message_replies_with_pattern(self): component = AnswerBuilder(pattern=r"Answer: (.*)") - replies = [ - ChatMessage( - content="Answer: AnswerString", - role=ChatRole.ASSISTANT, - name=None, - meta={ - "model": "gpt-4o-mini", - "index": 0, - "finish_reason": "stop", - "usage": {"prompt_tokens": 32, "completion_tokens": 153, "total_tokens": 185}, - }, - ) - ] - output = component.run(query="test query", replies=replies, meta=[{}]) - answers = output["answers"] - assert len(answers) == 1 - assert answers[0].data == "AnswerString" - assert answers[0].meta == { + + message_meta = { "model": "gpt-4o-mini", "index": 0, "finish_reason": "stop", "usage": {"prompt_tokens": 32, "completion_tokens": 153, "total_tokens": 185}, } + replies = [ChatMessage.from_assistant("Answer: AnswerString", meta=message_meta)] + + output = component.run(query="test query", replies=replies, meta=[{}]) + answers = output["answers"] + assert len(answers) == 1 + assert answers[0].data == "AnswerString" + assert answers[0].meta == message_meta assert answers[0].query == "test query" assert answers[0].documents == [] assert isinstance(answers[0], GeneratedAnswer) def test_run_with_chat_message_replies_with_documents(self): component = AnswerBuilder(reference_pattern="\\[(\\d+)\\]") - replies = [ - ChatMessage( - content="Answer: AnswerString[2]", - role=ChatRole.ASSISTANT, - name=None, - meta={ - "model": "gpt-4o-mini", - "index": 0, - "finish_reason": "stop", - "usage": {"prompt_tokens": 32, "completion_tokens": 153, "total_tokens": 185}, - }, - ) - ] + message_meta = { + "model": "gpt-4o-mini", + "index": 0, + "finish_reason": "stop", + "usage": {"prompt_tokens": 32, "completion_tokens": 153, "total_tokens": 185}, + } + replies = [ChatMessage.from_assistant("Answer: AnswerString[2]", meta=message_meta)] + output = component.run( query="test query", replies=replies, @@ -244,60 +221,40 @@ def test_run_with_chat_message_replies_with_documents(self): answers = output["answers"] assert len(answers) == 1 assert answers[0].data == "Answer: AnswerString[2]" - assert answers[0].meta == { - "model": "gpt-4o-mini", - "index": 0, - "finish_reason": "stop", - "usage": {"prompt_tokens": 32, "completion_tokens": 153, "total_tokens": 185}, - } + assert answers[0].meta == message_meta assert answers[0].query == "test query" assert len(answers[0].documents) == 1 assert answers[0].documents[0].content == "test doc 2" def test_run_with_chat_message_replies_with_pattern_set_at_runtime(self): component = AnswerBuilder(pattern="unused pattern") - replies = [ - ChatMessage( - content="Answer: AnswerString", - role=ChatRole.ASSISTANT, - name=None, - meta={ - "model": "gpt-4o-mini", - "index": 0, - "finish_reason": "stop", - "usage": {"prompt_tokens": 32, "completion_tokens": 153, "total_tokens": 185}, - }, - ) - ] - output = component.run(query="test query", replies=replies, meta=[{}], pattern=r"Answer: (.*)") - answers = output["answers"] - assert len(answers) == 1 - assert answers[0].data == "AnswerString" - assert answers[0].meta == { + message_meta = { "model": "gpt-4o-mini", "index": 0, "finish_reason": "stop", "usage": {"prompt_tokens": 32, "completion_tokens": 153, "total_tokens": 185}, } + replies = [ChatMessage.from_assistant("Answer: AnswerString", meta=message_meta)] + + output = component.run(query="test query", replies=replies, meta=[{}], pattern=r"Answer: (.*)") + answers = output["answers"] + assert len(answers) == 1 + assert answers[0].data == "AnswerString" + assert answers[0].meta == message_meta assert answers[0].query == "test query" assert answers[0].documents == [] assert isinstance(answers[0], GeneratedAnswer) def test_run_with_chat_message_replies_with_meta_set_at_run_time(self): component = AnswerBuilder() - replies = [ - ChatMessage( - content="AnswerString", - role=ChatRole.ASSISTANT, - name=None, - meta={ - "model": "gpt-4o-mini", - "index": 0, - "finish_reason": "stop", - "usage": {"prompt_tokens": 32, "completion_tokens": 153, "total_tokens": 185}, - }, - ) - ] + message_meta = { + "model": "gpt-4o-mini", + "index": 0, + "finish_reason": "stop", + "usage": {"prompt_tokens": 32, "completion_tokens": 153, "total_tokens": 185}, + } + replies = [ChatMessage.from_assistant("AnswerString", meta=message_meta)] + output = component.run(query="test query", replies=replies, meta=[{"test": "meta"}]) answers = output["answers"] assert len(answers) == 1 @@ -315,8 +272,9 @@ def test_run_with_chat_message_replies_with_meta_set_at_run_time(self): def test_run_with_chat_message_no_meta_with_meta_set_at_run_time(self): component = AnswerBuilder() - replies = [ChatMessage(content="AnswerString", role=ChatRole.ASSISTANT, name=None, meta={})] + replies = [ChatMessage.from_assistant("AnswerString")] output = component.run(query="test query", replies=replies, meta=[{"test": "meta"}]) + answers = output["answers"] assert len(answers) == 1 assert answers[0].data == "AnswerString" diff --git a/test/components/builders/test_chat_prompt_builder.py b/test/components/builders/test_chat_prompt_builder.py index cd5c2c62c8..7e42eedbcb 100644 --- a/test/components/builders/test_chat_prompt_builder.py +++ b/test/components/builders/test_chat_prompt_builder.py @@ -210,11 +210,14 @@ def test_run_with_meta(self): Test that the ChatPromptBuilder correctly handles meta data. It should render the message and copy the meta data from the original message. """ - m = ChatMessage(content="This is a {{ variable }}", role=ChatRole.USER, name=None, meta={"test": "test"}) + m = ChatMessage.from_user("This is a {{ variable }}") + m.meta["meta_field"] = "meta_value" builder = ChatPromptBuilder(template=[m]) res = builder.run(variable="test") - res_msg = ChatMessage(content="This is a test", role=ChatRole.USER, name=None, meta={"test": "test"}) - assert res == {"prompt": [res_msg]} + + expected_msg = ChatMessage.from_user("This is a test") + expected_msg.meta["meta_field"] = "meta_value" + assert res == {"prompt": [expected_msg]} def test_run_with_invalid_template(self): builder = ChatPromptBuilder() diff --git a/test/dataclasses/test_chat_message.py b/test/dataclasses/test_chat_message.py index 642ff4d557..165e403301 100644 --- a/test/dataclasses/test_chat_message.py +++ b/test/dataclasses/test_chat_message.py @@ -96,12 +96,12 @@ def test_to_dict(): def test_from_dict(): - assert ChatMessage.from_dict(data={"content": "text", "role": "user", "name": None}) == ChatMessage( - content="text", role=ChatRole("user"), name=None, meta={} + assert ChatMessage.from_dict(data={"content": "text", "role": "user", "name": None}) == ChatMessage.from_user( + "text" ) def test_from_dict_with_meta(): assert ChatMessage.from_dict( - data={"content": "text", "role": "user", "name": None, "meta": {"something": "something"}} - ) == ChatMessage(content="text", role=ChatRole("user"), name=None, meta={"something": "something"}) + data={"content": "text", "role": "assistant", "name": None, "meta": {"something": "something"}} + ) == ChatMessage.from_assistant("text", meta={"something": "something"})