simplification

deepset-ai · Dec 18, 2024 · 8f5a6da · 8f5a6da
1 parent 1ac7091
commit 8f5a6da
Showing 1 changed file with 3 additions and 7 deletions.
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
@@ -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])