chore: lint python code

langgenius · Sep 14, 2024 · e07b550 · e07b550
1 parent 421569a
commit e07b550
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/api/core/memory/token_buffer_memory.py b/api/core/memory/token_buffer_memory.py
@@ -55,7 +55,8 @@ def get_history_prompt_messages(
 
         messages = query.limit(message_limit).all()
 
-        # instead of all messages from the conversation, we only need to extract messages that belong to the thread of last message
+        # instead of all messages from the conversation, we only need to extract messages
+        # that belong to the thread of last message
         thread_messages = extract_thread_messages(messages)
         thread_messages.pop(0)
         messages = list(reversed(thread_messages))

diff --git a/api/core/prompt/utils/extract_thread_messages.py b/api/core/prompt/utils/extract_thread_messages.py
@@ -15,7 +15,7 @@ def extract_thread_messages(messages: list[dict]) -> list[dict]:
             thread_messages.append(message)
             next_message = message.parent_message_id
         else:
-            if message.id == next_message or next_message == UUID_NIL:
+            if next_message in {message.id, UUID_NIL}:
                 thread_messages.append(message)
                 next_message = message.parent_message_id