Merge pull request #224 from thenewboston-developers/feature/221-anth…

…ropic-support Better filtering messages for Anthropic
thenewboston-developers · Aug 9, 2024 · 4c542c5 · 4c542c5
2 parents bfe7387 + ae3c06d
commit 4c542c5
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 26 deletions.
diff --git a/thenewboston/discord/bot.py b/thenewboston/discord/bot.py
@@ -13,6 +13,9 @@
 
 from thenewboston.general.clients.llm import LLMClient, make_prompt_kwargs  # noqa: E402
 
+USER_ROLE = 'user'
+ASSISTANT_ROLE = 'assistant'
+
 logger = logging.getLogger(__name__)
 
 intents = discord.Intents.default()
@@ -32,7 +35,7 @@ def map_author_plaintext(author):
 
 
 def map_author_structured(author):
-    return 'assistant' if is_ia(author) else 'user'
+    return ASSISTANT_ROLE if is_ia(author) else USER_ROLE
 
 
 def messages_to_plaintext(messages):
@@ -56,6 +59,9 @@ def messages_to_structured(messages):
 
         prev_role = role
 
+    if structured_messages and structured_messages[0]['role'] != USER_ROLE:
+        structured_messages.pop(0)
+
     return structured_messages
 
 

diff --git a/thenewboston/discord/tests/test_bot.py b/thenewboston/discord/tests/test_bot.py
@@ -16,96 +16,107 @@ async def test_on_ready():
         await on_ready()
 
 
-@override_settings(IA_DISCORD_USER_ID=1234)
+@override_settings(IA_DISCORD_USER_ID=1)
 def test_messages_to_structured():
-    assert messages_to_structured([Message(author=Author(id=1234), content='hello')]) == [{
-        'role': 'assistant',
+    assert messages_to_structured([Message(author=Author(id=2), content='hello')]) == [{
+        'role': 'user',
         'content': [{
             'type': 'text',
             'text': 'hello'
         }]
     }]
     assert messages_to_structured([
-        Message(author=Author(id=1234), content='hello'),
-        Message(author=Author(id=1234), content='world')
+        Message(author=Author(id=2), content='hello'),
+        Message(author=Author(id=2), content='world')
     ]) == [{
-        'role': 'assistant',
+        'role': 'user',
         'content': [{
             'type': 'text',
             'text': 'hello\nworld'
         }]
     }]
     assert messages_to_structured([
-        Message(author=Author(id=1234), content='hello'),
-        Message(author=Author(id=10), content='world')
+        Message(author=Author(id=2), content='hello'),
+        Message(author=Author(id=1), content='world')
     ]) == [
         {
-            'role': 'assistant',
+            'role': 'user',
             'content': [{
                 'type': 'text',
                 'text': 'hello'
             }]
         },
         {
-            'role': 'user',
+            'role': 'assistant',
             'content': [{
                 'type': 'text',
                 'text': 'world'
             }]
         },
     ]
     assert messages_to_structured([
-        Message(author=Author(id=1234), content='hello'),
-        Message(author=Author(id=10), content='world'),
-        Message(author=Author(id=1234), content='bye')
+        Message(author=Author(id=2), content='hello'),
+        Message(author=Author(id=1), content='world'),
+        Message(author=Author(id=2), content='bye')
     ]) == [
         {
-            'role': 'assistant',
+            'role': 'user',
             'content': [{
                 'type': 'text',
                 'text': 'hello'
             }]
         },
         {
-            'role': 'user',
+            'role': 'assistant',
             'content': [{
                 'type': 'text',
                 'text': 'world'
             }]
         },
         {
-            'role': 'assistant',
+            'role': 'user',
             'content': [{
                 'type': 'text',
                 'text': 'bye'
             }]
         },
     ]
     assert messages_to_structured([
-        Message(author=Author(id=1234), content='hello'),
-        Message(author=Author(id=10), content='world'),
-        Message(author=Author(id=10), content='mine'),
-        Message(author=Author(id=1234), content='bye')
+        Message(author=Author(id=2), content='hello'),
+        Message(author=Author(id=1), content='world'),
+        Message(author=Author(id=2), content='mine'),
+        Message(author=Author(id=2), content='bye')
     ]) == [
         {
-            'role': 'assistant',
+            'role': 'user',
             'content': [{
                 'type': 'text',
                 'text': 'hello'
             }]
         },
         {
-            'role': 'user',
+            'role': 'assistant',
             'content': [{
                 'type': 'text',
-                'text': 'world\nmine'
+                'text': 'world'
             }]
         },
         {
-            'role': 'assistant',
+            'role': 'user',
             'content': [{
                 'type': 'text',
-                'text': 'bye'
+                'text': 'mine\nbye'
             }]
         },
     ]
+
+    assert messages_to_structured([
+        Message(author=Author(id=1), content='hello'),
+        Message(author=Author(id=2), content='world')
+    ]) == [{
+        'role': 'user',
+        'content': [{
+            'type': 'text',
+            'text': 'world'
+        }]
+    }]