Use ChatMessage class methods to build messages

deepset-ai · Nov 28, 2024 · 6abd658 · 6abd658
1 parent dd08e2e
commit 6abd658
Showing 1 changed file with 11 additions and 21 deletions.
diff --git a/test/components/builders/test_answer_builder.py b/test/components/builders/test_answer_builder.py
@@ -165,10 +165,8 @@ 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,
+            ChatMessage.from_assistant(
+                "Answer: AnswerString",
                 meta={
                     "model": "gpt-4o-mini",
                     "index": 0,
@@ -194,10 +192,8 @@ def test_run_with_chat_message_replies_without_pattern(self):
     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,
+            ChatMessage.from_assistant(
+                "Answer: AnswerString",
                 meta={
                     "model": "gpt-4o-mini",
                     "index": 0,
@@ -223,10 +219,8 @@ def test_run_with_chat_message_replies_with_pattern(self):
     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,
+            ChatMessage.from_assistant(
+                "Answer: AnswerString[2]",
                 meta={
                     "model": "gpt-4o-mini",
                     "index": 0,
@@ -257,10 +251,8 @@ def test_run_with_chat_message_replies_with_documents(self):
     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,
+            ChatMessage.from_assistant(
+                "Answer: AnswerString",
                 meta={
                     "model": "gpt-4o-mini",
                     "index": 0,
@@ -286,10 +278,8 @@ def test_run_with_chat_message_replies_with_pattern_set_at_runtime(self):
     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,
+            ChatMessage.from_assistant(
+                "AnswerString",
                 meta={
                     "model": "gpt-4o-mini",
                     "index": 0,
@@ -315,7 +305,7 @@ 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", meta={})]
         output = component.run(query="test query", replies=replies, meta=[{"test": "meta"}])
         answers = output["answers"]
         assert len(answers) == 1