From e1c0fc3b577f0ee603003a8b6546b1b482662080 Mon Sep 17 00:00:00 2001 From: Amna Mubashar Date: Fri, 30 Aug 2024 11:13:41 +0200 Subject: [PATCH] fix: chat roles for model responses in chat generators (#1030) --- .../amazon_bedrock/tests/test_chat_generator.py | 2 +- .../components/generators/google_ai/chat/gemini.py | 14 +++++++------- .../tests/generators/chat/test_chat_gemini.py | 3 ++- .../generators/google_vertex/chat/gemini.py | 14 +++++++------- .../ollama/examples/chat_generator_example.py | 2 +- 5 files changed, 18 insertions(+), 17 deletions(-) diff --git a/integrations/amazon_bedrock/tests/test_chat_generator.py b/integrations/amazon_bedrock/tests/test_chat_generator.py index ed0c27401..a455d2c93 100644 --- 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/integrations/google_ai/src/haystack_integrations/components/generators/google_ai/chat/gemini.py b/integrations/google_ai/src/haystack_integrations/components/generators/google_ai/chat/gemini.py index cf0005f39..e859a29fd 100644 --- a/integrations/google_ai/src/haystack_integrations/components/generators/google_ai/chat/gemini.py +++ b/integrations/google_ai/src/haystack_integrations/components/generators/google_ai/chat/gemini.py @@ -230,14 +230,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 @@ -250,13 +250,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: @@ -315,12 +315,12 @@ def _get_response(self, response_body: GenerateContentResponse) -> List[ChatMess for candidate in response_body.candidates: for part in candidate.content.parts: if part.text != "": - replies.append(ChatMessage.from_system(part.text)) + replies.append(ChatMessage.from_assistant(part.text)) elif part.function_call is not None: replies.append( ChatMessage( content=dict(part.function_call.args.items()), - role=ChatRole.SYSTEM, + role=ChatRole.ASSISTANT, name=part.function_call.name, ) ) @@ -343,4 +343,4 @@ def _get_stream_response( responses.append(content) combined_response = "".join(responses).lstrip() - return [ChatMessage.from_system(content=combined_response)] + return [ChatMessage.from_assistant(content=combined_response)] diff --git a/integrations/google_ai/tests/generators/chat/test_chat_gemini.py b/integrations/google_ai/tests/generators/chat/test_chat_gemini.py index 1a910b977..35ad8db14 100644 --- 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/integrations/google_vertex/src/haystack_integrations/components/generators/google_vertex/chat/gemini.py b/integrations/google_vertex/src/haystack_integrations/components/generators/google_vertex/chat/gemini.py index 893710121..e5ca1166d 100644 --- a/integrations/google_vertex/src/haystack_integrations/components/generators/google_vertex/chat/gemini.py +++ b/integrations/google_vertex/src/haystack_integrations/components/generators/google_vertex/chat/gemini.py @@ -161,12 +161,12 @@ 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) @@ -174,11 +174,11 @@ 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.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) @@ -233,12 +233,12 @@ def _get_response(self, response_body: GenerationResponse) -> List[ChatMessage]: for candidate in response_body.candidates: for part in candidate.content.parts: if part._raw_part.text != "": - replies.append(ChatMessage.from_system(part.text)) + replies.append(ChatMessage.from_assistant(part.text)) elif part.function_call is not None: replies.append( ChatMessage( content=dict(part.function_call.args.items()), - role=ChatRole.SYSTEM, + role=ChatRole.ASSISTANT, name=part.function_call.name, ) ) @@ -261,4 +261,4 @@ def _get_stream_response( responses.append(streaming_chunk.content) combined_response = "".join(responses).lstrip() - return [ChatMessage.from_system(content=combined_response)] + return [ChatMessage.from_assistant(content=combined_response)] diff --git a/integrations/ollama/examples/chat_generator_example.py b/integrations/ollama/examples/chat_generator_example.py index 834df78fb..2326ba708 100644 --- 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" ),