fix(agent): preserve complete model messages through tools and history

.github/workflows/knowledge-fs-ci.yml

@@ -87,6 +87,7 @@ jobs:
               - 'api/tests/unit_tests/models/test_agent_config_entities.py'
               - 'api/tests/unit_tests/controllers/console/workspace/rbac/test_roles.py'
               - 'api/tests/unit_tests/core/agent/test_base_agent_runner.py'
+              - 'api/tests/unit_tests/core/agent/test_message_replay.py'
               - 'api/tests/unit_tests/core/app/apps/agent_app/test_runtime_request_builder.py'
               - 'api/tests/unit_tests/core/workflow/nodes/agent_v2/test_runtime_request_builder.py'
               - 'api/tests/unit_tests/core/workflow/test_node_runtime.py'
@@ -263,6 +264,7 @@ jobs:
             api/tests/unit_tests/controllers/console/test_wraps.py
             api/tests/unit_tests/controllers/console/workspace/rbac/test_roles.py
             api/tests/unit_tests/core/agent/test_base_agent_runner.py
+            api/tests/unit_tests/core/agent/test_message_replay.py
             api/tests/unit_tests/core/app/apps/agent_app/test_runtime_request_builder.py
             api/tests/unit_tests/core/workflow/nodes/agent_v2/test_runtime_request_builder.py
             api/tests/unit_tests/core/workflow/test_node_runtime.py

api/.importlinter

@@ -29,6 +29,10 @@ layers =
 # error makes this list shrink whenever a legacy dependency is removed.
 unmatched_ignore_imports_alerting = error
 ignore_imports =
+    # Legacy agent callers delegate model history to the application owner during core migration.
+    core.agent.base_agent_runner -> services.agent.message_history
+    core.workflow.nodes.agent.message_transformer -> services.agent.message_history
+    core.workflow.nodes.agent.runtime_support -> services.agent.workflow_message_history
     core.app.app_config.easy_ui_based_app.dataset.manager -> services.knowledge.dataset_service
     core.app.apps.advanced_chat.app_generator -> controllers.console.app.workflow
     core.app.apps.advanced_chat.app_generator -> services.conversation_service

api/core/agent/base_agent_runner.py

@@ -4,6 +4,7 @@
 from decimal import Decimal
 from typing import TYPE_CHECKING, cast
 
+from pydantic import ValidationError
 from sqlalchemy import func, select
 from sqlalchemy.orm import Session
 
@@ -46,6 +47,7 @@
 from graphon.model_runtime.model_providers.base.large_language_model import LargeLanguageModel
 from models.enums import CreatorUserRole
 from models.model import Conversation, Message, MessageAgentThought, MessageFile, load_annotation_reply_config
+from services.agent.message_history import AgentMessageHistory
 
 if TYPE_CHECKING:
     from core.app.apps.workflow_app_runner import WorkflowRunDriver
@@ -294,15 +296,26 @@ def save_agent_thought(
         answer: str | None,
         messages_ids: list[str],
         llm_usage: LLMUsage | None = None,
+        model_messages: list[PromptMessage] | None = None,
     ):
         """
         Save agent thought
         """
-        stmt = select(MessageAgentThought).where(MessageAgentThought.id == agent_thought_id)
+        stmt = select(MessageAgentThought).where(
+            MessageAgentThought.id == agent_thought_id,
+            MessageAgentThought.message_id == self.message.id,
+        )
         agent_thought = db.session.scalar(stmt)
         if not agent_thought:
             raise ValueError("agent thought not found")
 
+        if model_messages is not None:
+            agent_thought.model_messages = AgentMessageHistory(
+                provider=self.model_config.provider,
+                model=self.model_config.model,
+                messages=model_messages,
+            ).model_dump_json()
+
         if thought:
             existing_thought = agent_thought.thought or ""
             agent_thought.thought = f"{existing_thought}{thought}"
@@ -385,7 +398,10 @@ def organize_agent_history(self, prompt_messages: list[PromptMessage], *, sessio
             (
                 session.execute(
                     select(Message)
-                    .where(Message.conversation_id == self.message.conversation_id)
+                    .where(
+                        Message.conversation_id == self.message.conversation_id,
+                        Message.app_id == self.app_config.app_id,
+                    )
                     .order_by(Message.created_at.desc())
                 )
             )
@@ -403,6 +419,17 @@ def organize_agent_history(self, prompt_messages: list[PromptMessage], *, sessio
             agent_thoughts = message.agent_thoughts_with_session(session=session)
             if agent_thoughts:
                 for agent_thought in agent_thoughts:
+                    if agent_thought.model_messages:
+                        try:
+                            history = AgentMessageHistory.model_validate_json(agent_thought.model_messages)
+                            restored = history.restore(
+                                provider=self.model_config.provider, model=self.model_config.model
+                            )
+                            if restored is not None:
+                                result.extend(restored)
+                                continue
+                        except ValidationError:
+                            logger.warning("Invalid model history for agent thought %s", agent_thought.id)
                     tool_names_raw = agent_thought.tool
                     if tool_names_raw:
                         tool_names = tool_names_raw.split(";")

api/core/agent/cot_agent_runner.py

@@ -25,9 +25,9 @@
     AssistantPromptMessage,
     PromptMessage,
     PromptMessageTool,
-    ToolPromptMessage,
     UserPromptMessage,
 )
+from graphon.model_runtime.model_providers.base.large_language_model import LLMResultAccumulator
 from models.model import Message
 
 logger = logging.getLogger(__name__)
@@ -129,9 +129,7 @@ def increase_usage(final_llm_usage_dict: dict[str, LLMUsage | None], usage: LLMU
                     QueueAgentThoughtEvent(agent_thought_id=agent_thought_id), PublishFrom.APPLICATION_MANAGER
                 )
 
-            # recalc llm max tokens
             prompt_messages = self._organize_prompt_messages()
-            self.recalc_llm_max_tokens(self.model_config, prompt_messages)
 
             # Release any setup/tool transaction before waiting on the provider stream.
             session.commit()
@@ -155,7 +153,18 @@ def increase_usage(final_llm_usage_dict: dict[str, LLMUsage | None], usage: LLMU
             )
 
             usage_dict: dict[str, LLMUsage | None] = {}
-            react_chunks = CotAgentOutputParser.handle_react_stream_output(chunks, usage_dict)
+            accumulator = LLMResultAccumulator()
+
+            def collect_response(
+                source: Generator[LLMResultChunk], target: LLMResultAccumulator
+            ) -> Generator[LLMResultChunk]:
+                for chunk in source:
+                    target.consume(chunk)
+                    yield chunk
+
+            react_chunks = CotAgentOutputParser.handle_react_stream_output(
+                collect_response(chunks, accumulator), usage_dict
+            )
             scratchpad = AgentScratchpadUnit(
                 agent_response="",
                 thought="",
@@ -190,6 +199,10 @@ def increase_usage(final_llm_usage_dict: dict[str, LLMUsage | None], usage: LLMU
                         delta=LLMResultChunkDelta(index=0, message=AssistantPromptMessage(content=chunk), usage=None),
                     )
 
+            thought_start = len(self._current_thoughts)
+            result = accumulator.to_result(model=model_instance.model_name, prompt_messages=prompt_messages)
+            self._current_thoughts.append(result.message)
+
             assert scratchpad.thought is not None
             scratchpad.thought = scratchpad.thought.strip() or "I am thinking about how to help you"
             self._agent_scratchpad.append(scratchpad)
@@ -216,6 +229,7 @@ def increase_usage(final_llm_usage_dict: dict[str, LLMUsage | None], usage: LLMU
                 answer=scratchpad.agent_response or "",
                 messages_ids=[],
                 llm_usage=usage_dict["usage"],
+                model_messages=self._current_thoughts[thought_start:],
             )
 
             if not scratchpad.is_final():
@@ -251,6 +265,11 @@ def increase_usage(final_llm_usage_dict: dict[str, LLMUsage | None], usage: LLMU
                     )
                     scratchpad.observation = tool_invoke_response
                     scratchpad.agent_response = tool_invoke_response
+                    self._current_thoughts.append(
+                        UserPromptMessage(
+                            content=f"Observation: {tool_invoke_response}\n\ncontinue", name="agent_observation"
+                        )
+                    )
 
                     self.save_agent_thought(
                         agent_thought_id=agent_thought_id,
@@ -262,6 +281,7 @@ def increase_usage(final_llm_usage_dict: dict[str, LLMUsage | None], usage: LLMU
                         answer=scratchpad.agent_response,
                         messages_ids=message_file_ids,
                         llm_usage=usage_dict["usage"],
+                        model_messages=self._current_thoughts[thought_start:],
                     )
 
                     self.queue_manager.publish(
@@ -388,6 +408,7 @@ def _init_react_state(self, query):
         """
         self._query = query
         self._agent_scratchpad = []
+        self._current_thoughts = []
         self._historic_prompt_messages = self._organize_historic_prompt_messages()
 
     @abstractmethod
@@ -419,53 +440,9 @@ def _organize_historic_prompt_messages(
         """
         organize historic prompt messages
         """
-        result: list[PromptMessage] = []
-        scratchpads: list[AgentScratchpadUnit] = []
-        current_scratchpad: AgentScratchpadUnit | None = None
-
-        for message in self.history_prompt_messages:
-            match message:
-                case AssistantPromptMessage():
-                    if not current_scratchpad:
-                        assert isinstance(message.content, str)
-                        current_scratchpad = AgentScratchpadUnit(
-                            agent_response=message.content,
-                            thought=message.content or "I am thinking about how to help you",
-                            action_str="",
-                            action=None,
-                            observation=None,
-                        )
-                        scratchpads.append(current_scratchpad)
-                    if message.tool_calls:
-                        try:
-                            current_scratchpad.action = AgentScratchpadUnit.Action(
-                                action_name=message.tool_calls[0].function.name,
-                                action_input=json.loads(message.tool_calls[0].function.arguments),
-                            )
-                            current_scratchpad.action_str = json.dumps(current_scratchpad.action.to_dict())
-                        except Exception:
-                            logger.exception("Failed to parse tool call from assistant message")
-                case ToolPromptMessage():
-                    if current_scratchpad:
-                        assert isinstance(message.content, str)
-                        current_scratchpad.observation = message.content
-                    else:
-                        raise NotImplementedError("expected str type")
-                case UserPromptMessage():
-                    if scratchpads:
-                        result.append(AssistantPromptMessage(content=self._format_assistant_message(scratchpads)))
-                        scratchpads = []
-                        current_scratchpad = None
-
-                    result.append(message)
-
-        if scratchpads:
-            result.append(AssistantPromptMessage(content=self._format_assistant_message(scratchpads)))
-
-        historic_prompts = AgentHistoryPromptTransform(
+        return AgentHistoryPromptTransform(
             model_config=self.model_config,
             prompt_messages=current_session_messages or [],
-            history_messages=result,
+            history_messages=self.history_prompt_messages,
             memory=self.memory,
         ).get_prompt()
-        return historic_prompts

api/core/agent/cot_chat_agent_runner.py

@@ -4,7 +4,6 @@
 from core.agent.cot_agent_runner import CotAgentRunner
 from graphon.file import file_manager
 from graphon.model_runtime.entities import (
-    AssistantPromptMessage,
     PromptMessage,
     SystemPromptMessage,
     TextPromptMessageContent,
@@ -75,43 +74,7 @@ def _organize_prompt_messages(self) -> list[PromptMessage]:
         # organize system prompt
         system_message = self._organize_system_prompt()
 
-        # organize current assistant messages
-        agent_scratchpad = self._agent_scratchpad
-        if not agent_scratchpad:
-            assistant_messages = []
-        else:
-            content = ""
-            for unit in agent_scratchpad:
-                if unit.is_final():
-                    content += f"Final Answer: {unit.agent_response}"
-                else:
-                    content += f"Thought: {unit.thought}\n\n"
-                    if unit.action_str:
-                        content += f"Action: {unit.action_str}\n\n"
-                    if unit.observation:
-                        content += f"Observation: {unit.observation}\n\n"
-
-            assistant_messages = [AssistantPromptMessage(content=content)]
-
-        # query messages
         query_messages = self._organize_user_query(self._query, [])
-
-        if assistant_messages:
-            # organize historic prompt messages
-            historic_messages = self._organize_historic_prompt_messages(
-                [system_message, *query_messages, *assistant_messages, UserPromptMessage(content="continue")]
-            )
-            messages = [
-                system_message,
-                *historic_messages,
-                *query_messages,
-                *assistant_messages,
-                UserPromptMessage(content="continue"),
-            ]
-        else:
-            # organize historic prompt messages
-            historic_messages = self._organize_historic_prompt_messages([system_message, *query_messages])
-            messages = [system_message, *historic_messages, *query_messages]
-
-        # join all messages
-        return messages
+        current_messages = [system_message, *query_messages, *self._current_thoughts]
+        historic_messages = self._organize_historic_prompt_messages(current_messages)
+        return [system_message, *historic_messages, *query_messages, *self._current_thoughts]