fix(openai): user and assistant roles should be alternating (yetone#859)

tmerse · Nov 16, 2024 · 9891b03 · 9891b03
1 parent ff85b9c
commit 9891b03
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 6 deletions.
diff --git a/lua/avante/providers/gemini.lua b/lua/avante/providers/gemini.lua
@@ -19,12 +19,15 @@ M.parse_messages = function(opts)
   vim.iter(opts.messages):each(function(message)
     local role = message.role
     if role == prev_role then
-      if role == "user" then
-        table.insert(contents, { role = "model", parts = {
-          { text = "Ok, I understand." },
-        } })
+      if role == M.role_map["user"] then
+        table.insert(
+          contents,
+          { role = M.role_map["assistant"], parts = {
+            { text = "Ok, I understand." },
+          } }
+        )
       else
-        table.insert(contents, { role = "user", parts = {
+        table.insert(contents, { role = M.role_map["user"], parts = {
           { text = "Ok" },
         } })
       end

diff --git a/lua/avante/providers/openai.lua b/lua/avante/providers/openai.lua
@@ -85,7 +85,23 @@ M.parse_messages = function(opts)
     messages[#messages].content = message_content
   end
 
-  return messages
+  local final_messages = {}
+  local prev_role = nil
+
+  vim.iter(messages):each(function(message)
+    local role = message.role
+    if role == prev_role then
+      if role == M.role_map["user"] then
+        table.insert(final_messages, { role = M.role_map["assistant"], content = "Ok, I understand." })
+      else
+        table.insert(final_messages, { role = M.role_map["user"], content = "Ok" })
+      end
+    end
+    prev_role = role
+    table.insert(final_messages, { role = M.role_map[role] or role, content = message.content })
+  end)
+
+  return final_messages
 end
 
 M.parse_response = function(data_stream, _, opts)