Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: chat roles for model responses in chat generators #1030

Merged
merged 5 commits into from
Aug 30, 2024
Merged

fix: chat roles for model responses in chat generators #1030

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
7483bc6
Replace system roles in responses
Amnah199 Aug 27, 2024
a6a648b
Add googleai fix
Amnah199 Aug 27, 2024
abc099a
Merge branch 'main' into fix-role-chat-gen
Amnah199 Aug 27, 2024
4c928ac
Update gemini.py
Amnah199 Aug 27, 2024
56106c3
Fix role handling
Amnah199 Aug 27, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion integrations/amazon_bedrock/tests/test_chat_generator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
14 changes: 7 additions & 7 deletions ...ations/google_ai/src/haystack_integrations/components/generators/google_ai/chat/gemini.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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}:
Amnah199 marked this conversation as resolved.
Show resolved Hide resolved
p = Part()
p.text = message.content
return p
Expand All @@ -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:
Expand Down Expand Up @@ -326,12 +326,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,
Amnah199 marked this conversation as resolved.
Show resolved Hide resolved
name=part.function_call.name,
)
)
Expand All @@ -354,4 +354,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)]
3 changes: 2 additions & 1 deletion integrations/google_ai/tests/generators/chat/test_chat_gemini.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
14 changes: 7 additions & 7 deletions ...oogle_vertex/src/haystack_integrations/components/generators/google_vertex/chat/gemini.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -236,12 +236,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,
)
)
Expand All @@ -264,4 +264,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)]
2 changes: 1 addition & 1 deletion integrations/ollama/examples/chat_generator_example.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
),
Expand Down