From 8f5a6da50c0cb8b6197663b52099179516bd824a Mon Sep 17 00:00:00 2001 From: anakin87 Date: Wed, 18 Dec 2024 17:09:22 +0100 Subject: [PATCH] simplification --- .../components/generators/google_ai/chat/gemini.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) 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 71e95eaed..69f168a6b 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 @@ -279,13 +279,9 @@ def _message_to_content(self, message: ChatMessage) -> Content: msg = f"Unsupported message role {message.role}" raise ValueError(msg) - role = "model" - if ( - message.is_from(ChatRole.USER) - or message.is_from(ChatRole.FUNCTION) - or ("TOOL" in ChatRole._member_names_ and message.is_from(ChatRole.TOOL)) - ): - role = "user" + role = "user" + if message.is_from(ChatRole.ASSISTANT) or message.is_from(ChatRole.SYSTEM): + role = "model" return Content(parts=[part], role=role) @component.output_types(replies=List[ChatMessage])