lib/gen_ai/api/format/gemini.rb in gen-ai-0.4.0.alpha.3 vs lib/gen_ai/api/format/gemini.rb in gen-ai-0.4.0

@@ -2,25 +2,32 @@
 
 module GenAI
   module Api
     module Format
       module Gemini
+        USER_ROLE = 'user'
+        ASSISTANT_ROLE = 'model'
+
         def format_messages(messages)
-          messages.map { |message| transform_message(message) }
+          messages.map { |message| transform_message(message.deep_symbolize_keys) }
         end
 
         def transform_message(message)
-          { role: role_for(message), parts: [text: message[:content]] }
+          if message.keys == %i[role content]
+            { role: role_for(message), parts: [text: message[:content]] }
+          else
+            message
+          end
         end
 
         def extract_completions(response)
           response['candidates'].map { |candidate| candidate.dig('content', 'parts', 0, 'text') }
         end
 
         private
 
         def role_for(message)
-          message[:role] == 'user' ? self.class::USER_ROLE : self.class::ASSISTANT_ROLE
+          message[:role] == 'user' ? USER_ROLE : ASSISTANT_ROLE
         end
       end
     end
   end
 end