From e07b550b8a31f5530b67869fe799a507e10890de Mon Sep 17 00:00:00 2001 From: xuzuodong Date: Sat, 14 Sep 2024 17:05:27 +0800 Subject: [PATCH] chore: lint python code --- api/core/memory/token_buffer_memory.py | 3 ++- api/core/prompt/utils/extract_thread_messages.py | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/api/core/memory/token_buffer_memory.py b/api/core/memory/token_buffer_memory.py index d55a7a761d03fb..60b36c50f00d38 100644 --- 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 index 3bd38dd30b2ba9..e8b626499fce19 100644 --- 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