feat: regenerate in Chat, agent and Chatflow app (#7661)

api/core/agent/base_agent_runner.py

@@ -32,6 +32,7 @@ from core.model_runtime.entities.message_entities import (
 from core.model_runtime.entities.model_entities import ModelFeature
 from core.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel
 from core.model_runtime.utils.encoders import jsonable_encoder
+from core.prompt.utils.extract_thread_messages import extract_thread_messages
 from core.tools.entities.tool_entities import (
     ToolParameter,
     ToolRuntimeVariablePool,
@@ -441,10 +442,12 @@ class BaseAgentRunner(AppRunner):
             .filter(
                 Message.conversation_id == self.message.conversation_id,
             )
-            .order_by(Message.created_at.asc())
+            .order_by(Message.created_at.desc())
             .all()
         )
 
+        messages = list(reversed(extract_thread_messages(messages)))
+
         for message in messages:
             if message.id == self.message.id:
                 continue

api/core/app/apps/advanced_chat/app_generator.py

@@ -121,6 +121,7 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
             inputs=conversation.inputs if conversation else self._get_cleaned_inputs(inputs, app_config),
             query=query,
             files=file_objs,
+            parent_message_id=args.get("parent_message_id"),
             user_id=user.id,
             stream=stream,
             invoke_from=invoke_from,

api/core/app/apps/agent_chat/app_generator.py

@@ -127,6 +127,7 @@ class AgentChatAppGenerator(MessageBasedAppGenerator):
             inputs=conversation.inputs if conversation else self._get_cleaned_inputs(inputs, app_config),
             query=query,
             files=file_objs,
+            parent_message_id=args.get("parent_message_id"),
             user_id=user.id,
             stream=stream,
             invoke_from=invoke_from,

api/core/app/apps/chat/app_generator.py

@@ -128,6 +128,7 @@ class ChatAppGenerator(MessageBasedAppGenerator):
             inputs=conversation.inputs if conversation else self._get_cleaned_inputs(inputs, app_config),
             query=query,
             files=file_objs,
+            parent_message_id=args.get("parent_message_id"),
             user_id=user.id,
             stream=stream,
             invoke_from=invoke_from,

api/core/app/apps/message_based_app_generator.py

@@ -218,6 +218,7 @@ class MessageBasedAppGenerator(BaseAppGenerator):
             answer_tokens=0,
             answer_unit_price=0,
             answer_price_unit=0,
+            parent_message_id=getattr(application_generate_entity, "parent_message_id", None),
             provider_response_latency=0,
             total_price=0,
             currency="USD",

api/core/app/entities/app_invoke_entities.py

@@ -122,6 +122,7 @@ class ChatAppGenerateEntity(EasyUIBasedAppGenerateEntity):
     """
 
     conversation_id: Optional[str] = None
+    parent_message_id: Optional[str] = None
 
 
 class CompletionAppGenerateEntity(EasyUIBasedAppGenerateEntity):
@@ -138,6 +139,7 @@ class AgentChatAppGenerateEntity(EasyUIBasedAppGenerateEntity):
     """
 
     conversation_id: Optional[str] = None
+    parent_message_id: Optional[str] = None
 
 
 class AdvancedChatAppGenerateEntity(AppGenerateEntity):
@@ -149,6 +151,7 @@ class AdvancedChatAppGenerateEntity(AppGenerateEntity):
     app_config: WorkflowUIBasedAppConfig
 
     conversation_id: Optional[str] = None
+    parent_message_id: Optional[str] = None
     query: str
 
     class SingleIterationRunEntity(BaseModel):

api/core/memory/token_buffer_memory.py

@@ -11,6 +11,7 @@ from core.model_runtime.entities.message_entities import (
     TextPromptMessageContent,
     UserPromptMessage,
 )
+from core.prompt.utils.extract_thread_messages import extract_thread_messages
 from extensions.ext_database import db
 from models.model import AppMode, Conversation, Message, MessageFile
 from models.workflow import WorkflowRun
@@ -33,8 +34,17 @@ class TokenBufferMemory:
 
         # fetch limited messages, and return reversed
         query = (
-            db.session.query(Message.id, Message.query, Message.answer, Message.created_at, Message.workflow_run_id)
-            .filter(Message.conversation_id == self.conversation.id, Message.answer != "")
+            db.session.query(
+                Message.id,
+                Message.query,
+                Message.answer,
+                Message.created_at,
+                Message.workflow_run_id,
+                Message.parent_message_id,
+            )
+            .filter(
+                Message.conversation_id == self.conversation.id,
+            )
             .order_by(Message.created_at.desc())
         )
 
@@ -45,7 +55,12 @@ class TokenBufferMemory:
 
         messages = query.limit(message_limit).all()
 
-        messages = list(reversed(messages))
+        # 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))
+
         message_file_parser = MessageFileParser(tenant_id=app_record.tenant_id, app_id=app_record.id)
         prompt_messages = []
         for message in messages:

api/core/prompt/utils/extract_thread_messages.py

@@ -0,0 +1,22 @@
+from constants import UUID_NIL
+
+
+def extract_thread_messages(messages: list[dict]) -> list[dict]:
+    thread_messages = []
+    next_message = None
+
+    for message in messages:
+        if not message.parent_message_id:
+            # If the message is regenerated and does not have a parent message, it is the start of a new thread
+            thread_messages.append(message)
+            break
+
+        if not next_message:
+            thread_messages.append(message)
+            next_message = message.parent_message_id
+        else:
+            if next_message in {message.id, UUID_NIL}:
+                thread_messages.append(message)
+                next_message = message.parent_message_id
+
+    return thread_messages