Fix role handling

deepset-ai · Aug 27, 2024 · 56106c3 · 56106c3
1 parent 4c928ac
commit 56106c3
Show file tree

Hide file tree

Showing 5 changed files with 13 additions and 12 deletions.
diff --git a/integrations/amazon_bedrock/tests/test_chat_generator.py b/integrations/amazon_bedrock/tests/test_chat_generator.py
@@ -200,7 +200,7 @@ def test_long_prompt_is_not_truncated_when_truncate_false(mock_boto3_session):
     """
     Test that a long prompt is not truncated and _ensure_token_limit is not called when truncate is set to False
     """
-    messages = [ChatMessage.from_system("What is the biggest city in United States?")]
+    messages = [ChatMessage.from_user("What is the biggest city in United States?")]
 
     # Our mock prompt is 8 tokens long, so it exceeds the total limit (8 prompt tokens + 3 generated tokens > 10 tokens)
     max_length_generated_text = 3

diff --git a/...ations/google_ai/src/haystack_integrations/components/generators/google_ai/chat/gemini.py b/...ations/google_ai/src/haystack_integrations/components/generators/google_ai/chat/gemini.py
@@ -241,14 +241,14 @@ def _convert_part(self, part: Union[str, ByteStream, Part]) -> Part:
             raise ValueError(msg)
 
     def _message_to_part(self, message: ChatMessage) -> Part:
-        if message.role == ChatRole.SYSTEM and message.name:
+        if message.role == ChatRole.ASSISTANT and message.name:
             p = Part()
             p.function_call.name = message.name
             p.function_call.args = {}
             for k, v in message.content.items():
                 p.function_call.args[k] = v
             return p
-        elif message.role == ChatRole.SYSTEM:
+        elif message.role in {ChatRole.SYSTEM, ChatRole.ASSISTANT}:
             p = Part()
             p.text = message.content
             return p
@@ -261,13 +261,13 @@ def _message_to_part(self, message: ChatMessage) -> Part:
             return self._convert_part(message.content)
 
     def _message_to_content(self, message: ChatMessage) -> Content:
-        if message.role == ChatRole.SYSTEM and message.name:
+        if message.role == ChatRole.ASSISTANT and message.name:
             part = Part()
             part.function_call.name = message.name
             part.function_call.args = {}
             for k, v in message.content.items():
                 part.function_call.args[k] = v
-        elif message.role == ChatRole.SYSTEM:
+        elif message.role in {ChatRole.SYSTEM, ChatRole.ASSISTANT}:
             part = Part()
             part.text = message.content
         elif message.role == ChatRole.FUNCTION:

diff --git a/integrations/google_ai/tests/generators/chat/test_chat_gemini.py b/integrations/google_ai/tests/generators/chat/test_chat_gemini.py
@@ -256,8 +256,9 @@ def get_current_weather(location: str, unit: str = "celsius"):  # noqa: ARG001
 def test_past_conversation():
     gemini_chat = GoogleAIGeminiChatGenerator(model="gemini-pro")
     messages = [
+        ChatMessage.from_system(content="You are a knowledageable mathematician."),
         ChatMessage.from_user(content="What is 2+2?"),
-        ChatMessage.from_system(content="It's an arithmetic operation."),
+        ChatMessage.from_assistant(content="It's an arithmetic operation."),
         ChatMessage.from_user(content="Yeah, but what's the result?"),
     ]
     res = gemini_chat.run(messages=messages)

diff --git a/...oogle_vertex/src/haystack_integrations/components/generators/google_vertex/chat/gemini.py b/...oogle_vertex/src/haystack_integrations/components/generators/google_vertex/chat/gemini.py
@@ -164,24 +164,24 @@ def _convert_part(self, part: Union[str, ByteStream, Part]) -> Part:
             raise ValueError(msg)
 
     def _message_to_part(self, message: ChatMessage) -> Part:
-        if message.role == ChatRole.SYSTEM and message.name:
+        if message.role == ChatRole.ASSISTANT and message.name:
             p = Part.from_dict({"function_call": {"name": message.name, "args": {}}})
             for k, v in message.content.items():
                 p.function_call.args[k] = v
             return p
-        elif message.role == ChatRole.SYSTEM:
+        elif message.role in {ChatRole.SYSTEM, ChatRole.ASSISTANT}:
             return Part.from_text(message.content)
         elif message.role == ChatRole.FUNCTION:
             return Part.from_function_response(name=message.name, response=message.content)
         elif message.role == ChatRole.USER:
             return self._convert_part(message.content)
 
     def _message_to_content(self, message: ChatMessage) -> Content:
-        if message.role == ChatRole.SYSTEM and message.name:
+        if message.role == ChatRole.ASSISTANT and message.name:
             part = Part.from_dict({"function_call": {"name": message.name, "args": {}}})
             for k, v in message.content.items():
                 part.function_call.args[k] = v
-        elif message.role == ChatRole.SYSTEM:
+        elif message.role in {ChatRole.SYSTEM, ChatRole.ASSISTANT}:
             part = Part.from_text(message.content)
         elif message.role == ChatRole.FUNCTION:
             part = Part.from_function_response(name=message.name, response=message.content)
@@ -241,7 +241,7 @@ def _get_response(self, response_body: GenerationResponse) -> List[ChatMessage]:
                     replies.append(
                         ChatMessage(
                             content=dict(part.function_call.args.items()),
-                            role=ChatRole.SYSTEM,
+                            role=ChatRole.ASSISTANT,
                             name=part.function_call.name,
                         )
                     )

diff --git a/integrations/ollama/examples/chat_generator_example.py b/integrations/ollama/examples/chat_generator_example.py
@@ -11,7 +11,7 @@
 
 messages = [
     ChatMessage.from_user("What's Natural Language Processing?"),
-    ChatMessage.from_system(
+    ChatMessage.from_assistant(
         "Natural Language Processing (NLP) is a field of computer science and artificial "
         "intelligence concerned with the interaction between computers and human language"
     ),