From 8b12bdf2b3196abea0e8db33d7c50fff48341cb9 Mon Sep 17 00:00:00 2001
From: Slavi Pantaleev <slavi@devture.com>
Date: Sun, 22 Sep 2024 06:39:00 +0000
Subject: [PATCH] Combine consecutive messages by the same user when talking to
 the Anthropic API

Fixes https://github.com/etkecc/baibot/issues/13
---
 src/agent/provider/anthropic/controller.rs |  9 +++
 src/conversation/llm/entity.rs             | 91 ++++++++++++++++++++++
 2 files changed, 100 insertions(+)

diff --git a/src/agent/provider/anthropic/controller.rs b/src/agent/provider/anthropic/controller.rs
index 3783a13..f305a06 100644
--- a/src/agent/provider/anthropic/controller.rs
+++ b/src/agent/provider/anthropic/controller.rs
@@ -111,6 +111,15 @@ impl ControllerTrait for Controller {
             })
         };
 
+        // Avoid the situation where multiple user or assistant messages are sent consecutively,
+        // to avoid errors like:
+        // > API error: Error response: error Api error: invalid_request_error messages: roles must alternate between "user" and "assistant", but found multiple "user" roles in a row
+        // as reported here: https://github.com/etkecc/baibot/issues/13
+        //
+        // As https://docs.anthropic.com/en/api/messages says:
+        // > Our models are trained to operate on alternating user and assistant conversational turns.
+        let conversation = conversation.combine_consecutive_messages();
+
         let mut conversation_messages = conversation.messages;
 
         if params.context_management_enabled {
diff --git a/src/conversation/llm/entity.rs b/src/conversation/llm/entity.rs
index 5a0fe73..671d406 100644
--- a/src/conversation/llm/entity.rs
+++ b/src/conversation/llm/entity.rs
@@ -15,3 +15,94 @@ pub struct Message {
 pub struct Conversation {
     pub messages: Vec<Message>,
 }
+
+impl Conversation {
+    /// Combine consecutive messages by the same author into a single message.
+    ///
+    /// Certain models (like Anthropic) cannot tolerate consecutive messages by the same author,
+    /// so combining them helps avoid issues.
+    /// See: https://github.com/etkecc/baibot/issues/13
+    pub fn combine_consecutive_messages(&self) -> Conversation {
+        // We'll likely get fewer messages, but let's reserve the maximum we expect.
+        let mut new_messages = Vec::with_capacity(self.messages.len());
+        let mut last_seen_author: Option<Author> = None;
+
+        for message in &self.messages {
+            let Some(last_seen_author_clone) = last_seen_author.clone() else {
+                last_seen_author = Some(message.author.clone());
+                new_messages.push(message.clone());
+                continue;
+            };
+
+            if message.author != last_seen_author_clone {
+                last_seen_author = Some(message.author.clone());
+                new_messages.push(message.clone());
+                continue;
+            }
+
+            new_messages.last_mut().unwrap().message_text.push('\n');
+            new_messages
+                .last_mut()
+                .unwrap()
+                .message_text
+                .push_str(&message.message_text);
+        }
+
+        Conversation {
+            messages: new_messages,
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn combine_consecutive_messages() {
+        let conversation = Conversation {
+            messages: vec![
+                Message {
+                    author: Author::User,
+                    message_text: "Hello".to_string(),
+                },
+                Message {
+                    author: Author::User,
+                    message_text: "How are you?".to_string(),
+                },
+                Message {
+                    author: Author::User,
+                    message_text: "I'm OK, btw.".to_string(),
+                },
+                Message {
+                    author: Author::Assistant,
+                    message_text: "Hi there!".to_string(),
+                },
+                Message {
+                    author: Author::Assistant,
+                    message_text: "I'm doing well, thank you.".to_string(),
+                },
+                Message {
+                    author: Author::User,
+                    message_text: "That's great!".to_string(),
+                },
+            ],
+        };
+
+        let conversation = conversation.combine_consecutive_messages();
+
+        assert_eq!(conversation.messages.len(), 3);
+        assert_eq!(conversation.messages[0].author, Author::User);
+        assert_eq!(
+            conversation.messages[0].message_text,
+            "Hello\nHow are you?\nI'm OK, btw."
+        );
+        assert_eq!(conversation.messages[1].author, Author::Assistant);
+        assert_eq!(
+            conversation.messages[1].message_text,
+            "Hi there!\nI'm doing well, thank you."
+        );
+        assert_eq!(conversation.messages[2].author, Author::User);
+        assert_eq!(conversation.messages[2].message_text, "That's great!");
+    }
+}