From 500f4b5aac345d287d4800b3bac6ef2e30ba05d7 Mon Sep 17 00:00:00 2001 From: WH-2099 Date: Wed, 21 Oct 2026 10:12:10 +0000 Subject: [PATCH 1/4] fix(agent): preserve complete model messages through tools and history --- api/.importlinter | 4 + api/core/agent/base_agent_runner.py | 31 +++- api/core/agent/cot_agent_runner.py | 77 +++------ api/core/agent/cot_chat_agent_runner.py | 43 +---- api/core/agent/fc_agent_runner.py | 129 ++------------ .../agent/output_parser/cot_output_parser.py | 4 +- .../output_parser/structured_output.py | 7 +- api/core/memory/token_buffer_memory.py | 17 +- .../prompt/agent_history_prompt_transform.py | 7 +- api/core/workflow/nodes/agent/agent_node.py | 2 + .../nodes/agent/message_transformer.py | 8 + .../workflow/nodes/agent/runtime_support.py | 12 +- .../e9c48a2701bd_add_agent_model_messages.py | 23 +++ api/models/model.py | 1 + api/pyproject.toml | 4 +- api/services/agent/message_history.py | 56 +++++++ .../agent/workflow_message_history.py | 63 +++++++ .../core/agent/test_base_agent_runner.py | 2 +- .../core/agent/test_cot_agent_runner.py | 14 +- .../core/agent/test_cot_chat_agent_runner.py | 52 +----- .../core/agent/test_fc_agent_runner.py | 157 ++++-------------- .../core/agent/test_message_replay.py | 122 ++++++++++++++ .../output_parser/test_structured_output.py | 30 ++-- .../test_agent_history_prompt_transform.py | 28 ++++ .../nodes/agent/test_message_history.py | 142 ++++++++++++++++ api/uv.lock | 9 +- .../user-manual/history-layer/index.md | 6 + dify-agent/pyproject.toml | 6 +- .../src/dify_agent/adapters/llm/model.py | 53 ++++-- .../dify_agent/adapters/llm/test_model.py | 94 +++++++++++ dify-agent/tests/local/test_packaging.py | 2 +- dify-agent/uv.lock | 56 +++---- 32 files changed, 800 insertions(+), 461 deletions(-) create mode 100644 api/migrations/versions/e9c48a2701bd_add_agent_model_messages.py create mode 100644 api/services/agent/message_history.py create mode 100644 api/services/agent/workflow_message_history.py create mode 100644 api/tests/unit_tests/core/agent/test_message_replay.py create mode 100644 api/tests/unit_tests/core/workflow/nodes/agent/test_message_history.py diff --git a/api/.importlinter b/api/.importlinter index 5f9ab46613b7d9..414ce2f0823902 100644 --- a/api/.importlinter +++ b/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 diff --git a/api/core/agent/base_agent_runner.py b/api/core/agent/base_agent_runner.py index db45e465131372..c38ed3cb25a386 100644 --- a/api/core/agent/base_agent_runner.py +++ b/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 @@ -43,6 +44,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 @@ -282,15 +284,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}" @@ -373,7 +386,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()) ) ) @@ -391,6 +407,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(";") diff --git a/api/core/agent/cot_agent_runner.py b/api/core/agent/cot_agent_runner.py index b508cc70497382..7da262857d85d0 100644 --- a/api/core/agent/cot_agent_runner.py +++ b/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 diff --git a/api/core/agent/cot_chat_agent_runner.py b/api/core/agent/cot_chat_agent_runner.py index 8d4a5931b821f7..c1b9a6a8a5d2ed 100644 --- a/api/core/agent/cot_chat_agent_runner.py +++ b/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] diff --git a/api/core/agent/fc_agent_runner.py b/api/core/agent/fc_agent_runner.py index 1b204608632436..3d03543885b72f 100644 --- a/api/core/agent/fc_agent_runner.py +++ b/api/core/agent/fc_agent_runner.py @@ -2,7 +2,6 @@ import logging import re from collections.abc import Generator -from copy import deepcopy from typing import Any from sqlalchemy import select @@ -29,13 +28,13 @@ LLMResultChunkDelta, LLMUsage, PromptMessage, - PromptMessageContentType, SystemPromptMessage, TextPromptMessageContent, ToolPromptMessage, UserPromptMessage, ) from graphon.model_runtime.entities.message_entities import ImagePromptMessageContent, PromptMessageContentUnionTypes +from graphon.model_runtime.model_providers.base.large_language_model import LLMResultAccumulator from models import UploadFile from models.model import Message @@ -158,9 +157,7 @@ def increase_usage(final_llm_usage_dict: dict[str, LLMUsage | None], usage: LLMU messages_ids=message_file_ids, ) - # 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() @@ -183,18 +180,9 @@ def increase_usage(final_llm_usage_dict: dict[str, LLMUsage | None], usage: LLMU request_metadata=request_metadata, ) - tool_calls: list[tuple[str, str, dict[str, Any]]] = [] - - # save full response - response = "" - - # save tool call names and inputs - tool_call_names = "" - tool_call_inputs = "" - current_llm_usage = None - if isinstance(chunks, Generator): + accumulator = LLMResultAccumulator() is_first_chunk = True for chunk in chunks: if is_first_chunk: @@ -202,88 +190,31 @@ def increase_usage(final_llm_usage_dict: dict[str, LLMUsage | None], usage: LLMU QueueAgentThoughtEvent(agent_thought_id=agent_thought_id), PublishFrom.APPLICATION_MANAGER ) is_first_chunk = False - # check if there is any tool call - if self.check_tool_calls(chunk): - function_call_state = True - tool_calls.extend(self.extract_tool_calls(chunk) or []) - tool_call_names = ";".join([tool_call[1] for tool_call in tool_calls]) - try: - tool_call_inputs = json.dumps( - {tool_call[1]: tool_call[2] for tool_call in tool_calls}, ensure_ascii=False - ) - except TypeError: - # fallback: force ASCII to handle non-serializable objects - tool_call_inputs = json.dumps({tool_call[1]: tool_call[2] for tool_call in tool_calls}) - - if chunk.delta.message and chunk.delta.message.content: - if isinstance(chunk.delta.message.content, list): - for content in chunk.delta.message.content: - response += content.data - else: - response += str(chunk.delta.message.content) - - if chunk.delta.usage: - increase_usage(llm_usage, chunk.delta.usage) - current_llm_usage = chunk.delta.usage - + accumulator.consume(chunk) yield chunk + result = accumulator.to_result(model=model_instance.model_name, prompt_messages=prompt_messages) else: result = chunks - # check if there is any tool call - if self.check_blocking_tool_calls(result): - function_call_state = True - tool_calls.extend(self.extract_blocking_tool_calls(result) or []) - tool_call_names = ";".join([tool_call[1] for tool_call in tool_calls]) - try: - tool_call_inputs = json.dumps( - {tool_call[1]: tool_call[2] for tool_call in tool_calls}, ensure_ascii=False - ) - except TypeError: - # fallback: force ASCII to handle non-serializable objects - tool_call_inputs = json.dumps({tool_call[1]: tool_call[2] for tool_call in tool_calls}) - - if result.usage: - increase_usage(llm_usage, result.usage) - current_llm_usage = result.usage - - if result.message and result.message.content: - if isinstance(result.message.content, list): - for content in result.message.content: - response += content.data - else: - response += result.message.content - - if not result.message.content: - result.message.content = "" - self.queue_manager.publish( QueueAgentThoughtEvent(agent_thought_id=agent_thought_id), PublishFrom.APPLICATION_MANAGER ) - yield LLMResultChunk( model=model_instance.model_name, prompt_messages=result.prompt_messages, system_fingerprint=result.system_fingerprint, - delta=LLMResultChunkDelta( - index=0, - message=result.message, - usage=result.usage, - ), + delta=LLMResultChunkDelta(index=0, message=result.message, usage=result.usage), ) - assistant_message = AssistantPromptMessage(content=response, tool_calls=[]) - if tool_calls: - assistant_message.tool_calls = [ - AssistantPromptMessage.ToolCall( - id=tool_call[0], - type="function", - function=AssistantPromptMessage.ToolCall.ToolCallFunction( - name=tool_call[1], arguments=json.dumps(tool_call[2], ensure_ascii=False) - ), - ) - for tool_call in tool_calls - ] - + if result.usage: + increase_usage(llm_usage, result.usage) + current_llm_usage = result.usage + assistant_message = result.message.model_copy(deep=True) + response = assistant_message.get_text_content() + tool_calls = self.extract_blocking_tool_calls(result) or [] + function_call_state = bool(tool_calls) + tool_call_names = ";".join(tool_call[1] for tool_call in tool_calls) + tool_call_inputs = json.dumps({call[1]: call[2] for call in tool_calls}, ensure_ascii=False) + thought_start = len(self._current_thoughts) self._current_thoughts.append(assistant_message) # save thought @@ -297,6 +228,7 @@ def increase_usage(final_llm_usage_dict: dict[str, LLMUsage | None], usage: LLMU answer=response, messages_ids=[], llm_usage=current_llm_usage, + model_messages=self._current_thoughts[thought_start:], ) self.queue_manager.publish( QueueAgentThoughtEvent(agent_thought_id=agent_thought_id), PublishFrom.APPLICATION_MANAGER @@ -396,6 +328,7 @@ def increase_usage(final_llm_usage_dict: dict[str, LLMUsage | None], usage: LLMU }, answer="", messages_ids=message_file_ids, + model_messages=self._current_thoughts[thought_start:], ) self.queue_manager.publish( QueueAgentThoughtEvent(agent_thought_id=agent_thought_id), PublishFrom.APPLICATION_MANAGER @@ -527,31 +460,6 @@ def _organize_user_query(self, query: str, prompt_messages: list[PromptMessage]) return prompt_messages - def _clear_user_prompt_image_messages(self, prompt_messages: list[PromptMessage]) -> list[PromptMessage]: - """ - As for now, gpt supports both fc and vision at the first iteration. - We need to remove the image messages from the prompt messages at the first iteration. - """ - prompt_messages = deepcopy(prompt_messages) - - for prompt_message in prompt_messages: - if isinstance(prompt_message, UserPromptMessage): - if prompt_message.name == _KNOWLEDGE_RETRIEVAL_PROMPT_NAME: - continue - if isinstance(prompt_message.content, list): - prompt_message.content = "\n".join( - [ - content.data - if content.type == PromptMessageContentType.TEXT - else "[image]" - if content.type == PromptMessageContentType.IMAGE - else "[file]" - for content in prompt_message.content - ] - ) - - return prompt_messages - def _organize_prompt_messages(self): prompt_template = self.app_config.prompt_template.simple_prompt_template or "" self.history_prompt_messages = self._init_system_message(prompt_template, self.history_prompt_messages) @@ -565,7 +473,4 @@ def _organize_prompt_messages(self): ).get_prompt() prompt_messages = [*self.history_prompt_messages, *query_prompt_messages, *self._current_thoughts] - if len(self._current_thoughts) != 0: - # clear messages after the first iteration - prompt_messages = self._clear_user_prompt_image_messages(prompt_messages) return prompt_messages diff --git a/api/core/agent/output_parser/cot_output_parser.py b/api/core/agent/output_parser/cot_output_parser.py index 75da9133b5f335..ee8c4a23d73dba 100644 --- a/api/core/agent/output_parser/cot_output_parser.py +++ b/api/core/agent/output_parser/cot_output_parser.py @@ -73,9 +73,7 @@ def extra_json_from_code_block(code_block) -> list[list | dict]: for response in llm_response: if response.delta.usage: usage_dict["usage"] = response.delta.usage - response_content = response.delta.message.content - if not isinstance(response_content, str): - continue + response_content = response.delta.message.get_text_content() # stream index = 0 diff --git a/api/core/llm_generator/output_parser/structured_output.py b/api/core/llm_generator/output_parser/structured_output.py index 8205c184a31e52..b74498b48cdcc6 100644 --- a/api/core/llm_generator/output_parser/structured_output.py +++ b/api/core/llm_generator/output_parser/structured_output.py @@ -147,13 +147,8 @@ def invoke_llm_with_structured_output( ) if isinstance(llm_result, LLMResult): - if not isinstance(llm_result.message.content, str): - raise OutputParserError( - f"Failed to parse structured output, LLM result is not a string: {llm_result.message.content}" - ) - return LLMResultWithStructuredOutput( - structured_output=_parse_structured_output(llm_result.message.content), + structured_output=_parse_structured_output(llm_result.message.get_text_content()), model=llm_result.model, message=llm_result.message, usage=llm_result.usage, diff --git a/api/core/memory/token_buffer_memory.py b/api/core/memory/token_buffer_memory.py index bb0d50357ffca9..cc7c0a6e06f957 100644 --- a/api/core/memory/token_buffer_memory.py +++ b/api/core/memory/token_buffer_memory.py @@ -1,5 +1,5 @@ from collections import defaultdict -from collections.abc import Sequence +from collections.abc import Callable, Sequence from dataclasses import dataclass from sqlalchemy import select @@ -87,10 +87,13 @@ class PreparedHistory: perform storage/network I/O, so callers should first close the history session. """ - prompts: tuple[HistoryPrompt, ...] + prompts: tuple[HistoryPrompt | PromptMessage, ...] + + def to_prompt_messages(self) -> list[PromptMessage]: + return [prompt.to_prompt_message() if isinstance(prompt, HistoryPrompt) else prompt for prompt in self.prompts] def get_prompt_messages(self, *, model_instance: ModelInstance, max_token_limit: int) -> Sequence[PromptMessage]: - prompt_messages = [prompt.to_prompt_message() for prompt in self.prompts] + prompt_messages = self.to_prompt_messages() if not prompt_messages: return [] @@ -207,6 +210,7 @@ def load_history( session: Session, message_limit: int | None, workflow_run_repo: APIWorkflowRunRepository | None = None, + message_loader: Callable[[Message], Sequence[PromptMessage] | None] | None = None, ) -> PreparedHistory: """Read ordered history and attachment configuration without provider or file I/O. @@ -258,8 +262,13 @@ def load_history( ).all(): assistant_files_by_message[message_file.message_id].append(message_file) - prompts: list[HistoryPrompt] = [] + prompts: list[HistoryPrompt | PromptMessage] = [] for message in messages: + if message_loader is not None: + restored = message_loader(message) + if restored is not None: + prompts.extend(restored) + continue user_files = user_files_by_message.get(message.id, []) assistant_files = assistant_files_by_message.get(message.id, []) if ( diff --git a/api/core/prompt/agent_history_prompt_transform.py b/api/core/prompt/agent_history_prompt_transform.py index 7c6280fe937b81..5585aa1df9cd58 100644 --- a/api/core/prompt/agent_history_prompt_transform.py +++ b/api/core/prompt/agent_history_prompt_transform.py @@ -62,8 +62,11 @@ def get_prompt(self) -> list[PromptMessage]: continue prompt_messages.append(prompt_message) num_prompt += 1 - # a message is start with UserPromptMessage - if isinstance(prompt_message, UserPromptMessage): + # Tool observations belong to the original user turn. + if isinstance(prompt_message, UserPromptMessage) and prompt_message.name not in { + "agent_observation", + "knowledge_retrieval", + }: curr_message_tokens = model_type_instance.get_num_tokens( self.model_config.model, self.model_config.credentials, diff --git a/api/core/workflow/nodes/agent/agent_node.py b/api/core/workflow/nodes/agent/agent_node.py index d5036f09f4c12d..d5087eda46f198 100644 --- a/api/core/workflow/nodes/agent/agent_node.py +++ b/api/core/workflow/nodes/agent/agent_node.py @@ -121,6 +121,7 @@ def _run(self) -> Generator[NodeEventPayload]: agent_parameters = strategy.get_parameters() parameters = self._runtime_support.build_parameters( + node_id=self._node_id, agent_parameters=agent_parameters, variable_pool=self.runtime_state.variable_pool, node_data=self.node_data, @@ -131,6 +132,7 @@ def _run(self) -> Generator[NodeEventPayload]: invoke_from=dify_ctx.invoke_from, ) parameters_for_log = self._runtime_support.build_parameters( + node_id=self._node_id, agent_parameters=agent_parameters, variable_pool=self.runtime_state.variable_pool, node_data=self.node_data, diff --git a/api/core/workflow/nodes/agent/message_transformer.py b/api/core/workflow/nodes/agent/message_transformer.py index de64593159f4bd..08c3204229bb32 100644 --- a/api/core/workflow/nodes/agent/message_transformer.py +++ b/api/core/workflow/nodes/agent/message_transformer.py @@ -23,6 +23,7 @@ ) from graphon.variables.segments import ArrayFileSegment from models import ToolFile +from services.agent.message_history import AgentMessageHistory from services.tools.builtin_tools_manage_service import BuiltinToolManageService from .events import AgentLogEvent @@ -63,6 +64,7 @@ def transform( agent_execution_metadata: Mapping[WorkflowNodeExecutionMetadataKey, Any] = {} llm_usage = LLMUsage.empty_usage() variables: dict[str, Any] = {} + process_data: dict[str, Any] = {} for message in message_stream: if message.type in { @@ -136,6 +138,11 @@ def transform( assert isinstance(message.message, ToolInvokeMessage.JsonMessage) if node_type == BuiltinNodeTypes.AGENT: if isinstance(message.message.json_object, dict): + model_history = message.message.json_object.pop("model_history", None) + if model_history is not None: + process_data["model_history"] = AgentMessageHistory.model_validate( + model_history + ).model_dump(mode="json") msg_metadata: dict[str, Any] = message.message.json_object.pop("execution_metadata", {}) llm_usage = LLMUsage.from_metadata(cast(LLMUsageMetadata, msg_metadata)) agent_execution_metadata = { @@ -306,6 +313,7 @@ def transform( WorkflowNodeExecutionMetadataKey.AGENT_LOG: agent_logs, }, inputs=parameters_for_log, + process_data=process_data, llm_usage=llm_usage, ) ) diff --git a/api/core/workflow/nodes/agent/runtime_support.py b/api/core/workflow/nodes/agent/runtime_support.py index 6ffccd70777e33..24b7f129a74159 100644 --- a/api/core/workflow/nodes/agent/runtime_support.py +++ b/api/core/workflow/nodes/agent/runtime_support.py @@ -23,6 +23,7 @@ from graphon.runtime import VariablePool from graphon.variables.template_resolution import convert_template from models.model import Conversation +from services.agent.workflow_message_history import get_workflow_agent_history from .entities import AgentNodeData, AgentOldVersionModelFeatures, ParamsAutoGenerated from .exceptions import AgentInputTypeError, AgentVariableNotFoundError @@ -40,6 +41,7 @@ def build_parameters( tenant_id: str, user_id: str, app_id: str, + node_id: str, invoke_from: Any, for_log: bool = False, ) -> dict[str, Any]: @@ -189,8 +191,14 @@ def build_parameters( model_instance=model_instance, ) if memory: - prompt_messages = memory.get_history_prompt_messages( - message_limit=node_data.memory.window.size or None + prompt_messages = get_workflow_agent_history( + memory=memory, + tenant_id=tenant_id, + app_id=app_id, + node_id=node_id, + provider=value.get("provider", ""), + model=value.get("model", ""), + message_limit=node_data.memory.window.size or None, ) history_prompt_messages = [ prompt_message.model_dump(mode="json") for prompt_message in prompt_messages diff --git a/api/migrations/versions/e9c48a2701bd_add_agent_model_messages.py b/api/migrations/versions/e9c48a2701bd_add_agent_model_messages.py new file mode 100644 index 00000000000000..055a7dd99ef77b --- /dev/null +++ b/api/migrations/versions/e9c48a2701bd_add_agent_model_messages.py @@ -0,0 +1,23 @@ +"""Persist complete agent model messages. + +Revision ID: e9c48a2701bd +Revises: a8c3e7f4b912 +""" + +import sqlalchemy as sa +from alembic import op + +import models.types + +revision = "e9c48a2701bd" +down_revision = "a8c3e7f4b912" +branch_labels = None +depends_on = None + + +def upgrade(): + op.add_column("message_agent_thoughts", sa.Column("model_messages", models.types.LongText(), nullable=True)) + + +def downgrade(): + op.drop_column("message_agent_thoughts", "model_messages") diff --git a/api/models/model.py b/api/models/model.py index 11c7ec4128cc83..a16c7a997ef41e 100644 --- a/api/models/model.py +++ b/api/models/model.py @@ -2451,6 +2451,7 @@ class MessageAgentThought(TypeBase): # plugin_id = mapped_column(StringUUID, nullable=True) ## for future design tool_process_data: Mapped[str | None] = mapped_column(LongText, nullable=True, default=None) message: Mapped[str | None] = mapped_column(LongText, nullable=True, default=None) + model_messages: Mapped[str | None] = mapped_column(LongText, nullable=True, default=None) message_token: Mapped[int | None] = mapped_column(sa.Integer, nullable=True, default=None) message_unit_price: Mapped[Decimal | None] = mapped_column(sa.Numeric, nullable=True, default=None) message_price_unit: Mapped[Decimal] = mapped_column(sa.Numeric(10, 7), nullable=False, default=Decimal("0.001")) diff --git a/api/pyproject.toml b/api/pyproject.toml index e13bd2a2dcba1a..33b61486db1cfb 100644 --- a/api/pyproject.toml +++ b/api/pyproject.toml @@ -63,8 +63,8 @@ exclude = ["providers/vdb/__pycache__", "providers/trace/__pycache__"] [tool.uv.sources] dify-agent = { path = "../dify-agent", editable = true } flask-restx = { git = "https://github.com/asukaminato0721/flask-restx", rev = "27758e26f8f740d7525d5039c51a9e524b6e2b68" } -# shortcut: pin the 0.8.0 container-type fix until a release includes graphon#344. -graphon = { git = "https://github.com/langgenius/graphon", rev = "4fa22141774dff80e90e4dfa98ea3e6fbf0421bb" } +# shortcut: pin the joint replay candidate until Graphon releases incremental accumulation and container output fixes. +graphon = { git = "https://github.com/langgenius/graphon", rev = "bfeca97726f91eab257e513357b72d456de851b2" } dify-vdb-alibabacloud-mysql = { workspace = true } dify-vdb-analyticdb = { workspace = true } dify-vdb-baidu = { workspace = true } diff --git a/api/services/agent/message_history.py b/api/services/agent/message_history.py new file mode 100644 index 00000000000000..e2bfa097d5117b --- /dev/null +++ b/api/services/agent/message_history.py @@ -0,0 +1,56 @@ +from collections.abc import Mapping +from typing import Any + +from pydantic import BaseModel, SerializeAsAny, field_validator + +from graphon.model_runtime.entities.message_entities import ( + AssistantPromptMessage, + PromptMessage, + SystemPromptMessage, + ToolPromptMessage, + UserPromptMessage, +) + + +class AgentMessageHistory(BaseModel): + provider: str + model: str + messages: list[SerializeAsAny[PromptMessage]] + + @field_validator("messages", mode="before") + @classmethod + def parse_messages(cls, values: Any) -> list[PromptMessage]: + message_types = { + "system": SystemPromptMessage, + "user": UserPromptMessage, + "assistant": AssistantPromptMessage, + "tool": ToolPromptMessage, + } + if not isinstance(values, list): + raise ValueError("Agent history messages must be a list") + messages = [] + for value in values: + if isinstance(value, PromptMessage): + messages.append(value) + elif isinstance(value, Mapping) and value.get("role") in message_types: + messages.append(message_types[value["role"]].model_validate(value)) + else: + raise ValueError("Invalid agent history message role") + return messages + + def restore(self, *, provider: str, model: str) -> list[PromptMessage] | None: + if (self.provider, self.model) != (provider, model): + return None + messages = list(self.messages) + pending: dict[str, str] = {} + for message in messages: + if isinstance(message, AssistantPromptMessage): + pending.update((call.id, call.function.name) for call in message.tool_calls) + elif isinstance(message, ToolPromptMessage): + pending.pop(message.tool_call_id, None) + # A stopped run can persist the model response before its tools finish. + messages.extend( + ToolPromptMessage(tool_call_id=call_id, name=name, content="Tool execution was interrupted.") + for call_id, name in pending.items() + ) + return messages diff --git a/api/services/agent/workflow_message_history.py b/api/services/agent/workflow_message_history.py new file mode 100644 index 00000000000000..813b128c2f915d --- /dev/null +++ b/api/services/agent/workflow_message_history.py @@ -0,0 +1,63 @@ +import logging + +from pydantic import ValidationError +from sqlalchemy import select +from sqlalchemy.orm import sessionmaker + +from core.memory.token_buffer_memory import TokenBufferMemory +from extensions.ext_database import db +from graphon.model_runtime.entities.message_entities import PromptMessage +from models.model import App, Message +from repositories.factory import DifyAPIRepositoryFactory +from services.agent.message_history import AgentMessageHistory + +logger = logging.getLogger(__name__) + + +def get_workflow_agent_history( + *, + memory: TokenBufferMemory, + tenant_id: str, + app_id: str, + node_id: str, + provider: str, + model: str, + message_limit: int | None, +) -> list[PromptMessage]: + session_maker = sessionmaker(bind=db.engine, expire_on_commit=False) + repository = DifyAPIRepositoryFactory.create_api_workflow_node_execution_repository(session_maker) + + def load_message(message: Message) -> list[PromptMessage] | None: + if message.app_id != app_id: + return [] + if not message.workflow_run_id: + return None + executions = repository.get_executions_by_workflow_run(tenant_id, app_id, message.workflow_run_id) + restored: list[PromptMessage] = [] + for execution in executions: + if execution.node_id != node_id: + continue + process_data = repository.load_full_process_data(execution) + if not process_data or "model_history" not in process_data: + continue + try: + history = AgentMessageHistory.model_validate(process_data["model_history"]) + messages = history.restore(provider=provider, model=model) + if messages is not None: + restored.extend(messages) + except ValidationError: + logger.warning("Invalid agent model history for execution %s", execution.id) + return restored or None + + with session_maker() as session: + app = session.scalar(select(App).where(App.id == app_id, App.tenant_id == tenant_id)) + if app is None or memory.conversation.app_id != app_id: + return [] + history = TokenBufferMemory.load_history( + conversation=memory.conversation, + app_record=app, + session=session, + message_limit=message_limit, + message_loader=load_message, + ) + return history.to_prompt_messages() diff --git a/api/tests/unit_tests/core/agent/test_base_agent_runner.py b/api/tests/unit_tests/core/agent/test_base_agent_runner.py index 476d743214878b..aaaea4177929c1 100644 --- a/api/tests/unit_tests/core/agent/test_base_agent_runner.py +++ b/api/tests/unit_tests/core/agent/test_base_agent_runner.py @@ -319,7 +319,7 @@ def test_create_agent_thought_persists_row( def _persist_thought(session: Session, *, tool: str = "tool1;tool2") -> MessageAgentThought: - thought = _thought(thought_id="thought-1", tool=tool, thought="") + thought = _thought(thought_id="thought-1", message_id="msg_current", tool=tool, thought="") session.add(thought) session.commit() return thought diff --git a/api/tests/unit_tests/core/agent/test_cot_agent_runner.py b/api/tests/unit_tests/core/agent/test_cot_agent_runner.py index 771b819fa02119..d50bb0d028fa0e 100644 --- a/api/tests/unit_tests/core/agent/test_cot_agent_runner.py +++ b/api/tests/unit_tests/core/agent/test_cot_agent_runner.py @@ -14,7 +14,8 @@ from core.agent.errors import AgentMaxIterationError from core.app.entities.app_invoke_entities import CreditUsageCreatedBy from core.credit_usage import CreditUsageAppType -from graphon.model_runtime.entities.llm_entities import LLMUsage +from graphon.model_runtime.entities.llm_entities import LLMResultChunk, LLMResultChunkDelta, LLMUsage +from graphon.model_runtime.entities.message_entities import AssistantPromptMessage from libs.datetime_utils import naive_utc_now from models.enums import ConversationFromSource, MessageStatus from models.model import AppMode, Conversation, Message @@ -396,12 +397,14 @@ def record_detach(_session: Session, instance: object) -> None: def provider_chunks(): events.append("first-chunk") - yield "chunk" + yield LLMResultChunk( + model="m", delta=LLMResultChunkDelta(index=0, message=AssistantPromptMessage(content="chunk")) + ) runner.model_instance.invoke_llm.return_value = provider_chunks() mocker.patch( "core.agent.cot_agent_runner.CotAgentOutputParser.handle_react_stream_output", - side_effect=lambda chunks, _usage: list(chunks), + side_effect=lambda chunks, _usage: [chunk.delta.message.get_text_content() for chunk in chunks], ) results = list(runner.run(session, message, "query", {})) @@ -544,13 +547,12 @@ def test_user_message_flushes_scratchpad(self, runner: DummyRunner, mocker: Mock result = runner._organize_historic_prompt_messages([]) assert result == ["final"] - def test_tool_message_without_scratchpad_raises(self, runner: DummyRunner): + def test_history_without_memory_is_omitted(self, runner: DummyRunner): from graphon.model_runtime.entities.message_entities import ToolPromptMessage runner.history_prompt_messages = [ToolPromptMessage(content="obs", tool_call_id="1")] - with pytest.raises(NotImplementedError): - runner._organize_historic_prompt_messages([]) + assert runner._organize_historic_prompt_messages([]) == [] def test_agent_history_transform_invocation(self, runner: DummyRunner, mocker: MockerFixture): mock_transform = MagicMock() diff --git a/api/tests/unit_tests/core/agent/test_cot_chat_agent_runner.py b/api/tests/unit_tests/core/agent/test_cot_chat_agent_runner.py index 8e7093fd123a1d..38ab35ad271588 100644 --- a/api/tests/unit_tests/core/agent/test_cot_chat_agent_runner.py +++ b/api/tests/unit_tests/core/agent/test_cot_chat_agent_runner.py @@ -49,6 +49,7 @@ def runner(): runner._prompt_messages_tools = [DummyTool("tool1"), DummyTool("tool2")] runner._query = "user query" runner._agent_scratchpad = [] + runner._current_thoughts = [] runner.files = [] runner.application_generate_entity = DummyGenerateEntity() runner._organize_historic_prompt_messages = MagicMock(return_value=["historic"]) @@ -165,52 +166,15 @@ def test_no_scratchpad(self, runner, mocker: MockerFixture): assert "query" in result runner._organize_historic_prompt_messages.assert_called_once() - def test_with_final_scratchpad(self, runner, mocker: MockerFixture): - runner.app_config = DummyAppConfig(DummyAgentConfig(DummyPrompt("{{instruction}}"))) - runner._organize_system_prompt = MagicMock(return_value="system") - runner._organize_user_query = MagicMock(return_value=["query"]) + def test_preserves_model_messages_between_tool_rounds(self, runner): + from graphon.model_runtime.entities.message_entities import AssistantPromptMessage, UserPromptMessage - unit = DummyUnit(final=True, agent_response="done") - runner._agent_scratchpad = [unit] - - result = runner._organize_prompt_messages() - assistant_msgs = [m for m in result if hasattr(m, "content")] - combined = "".join([m.content for m in assistant_msgs if isinstance(m.content, str)]) - assert "Final Answer: done" in combined - - def test_with_thought_action_observation(self, runner, mocker: MockerFixture): - runner.app_config = DummyAppConfig(DummyAgentConfig(DummyPrompt("{{instruction}}"))) runner._organize_system_prompt = MagicMock(return_value="system") runner._organize_user_query = MagicMock(return_value=["query"]) - - unit = DummyUnit( - final=False, - thought="thinking", - action_str="action", - observation="observe", - ) - runner._agent_scratchpad = [unit] - - result = runner._organize_prompt_messages() - assistant_msgs = [m for m in result if hasattr(m, "content")] - combined = "".join([m.content for m in assistant_msgs if isinstance(m.content, str)]) - assert "Thought: thinking" in combined - assert "Action: action" in combined - assert "Observation: observe" in combined - - def test_multiple_units_mixed(self, runner, mocker: MockerFixture): - runner.app_config = DummyAppConfig(DummyAgentConfig(DummyPrompt("{{instruction}}"))) - runner._organize_system_prompt = MagicMock(return_value="system") - runner._organize_user_query = MagicMock(return_value=["query"]) - - units = [ - DummyUnit(final=False, thought="t1"), - DummyUnit(final=True, agent_response="done"), + runner._current_thoughts = [ + AssistantPromptMessage(content=[TextPromptMessageContent(data="action", opaque_body=0)], opaque_body=False), + UserPromptMessage(content="Observation: result\n\ncontinue"), + AssistantPromptMessage(content="done", opaque_body={"signature": "signed"}), ] - runner._agent_scratchpad = units - result = runner._organize_prompt_messages() - assistant_msgs = [m for m in result if hasattr(m, "content")] - combined = "".join([m.content for m in assistant_msgs if isinstance(m.content, str)]) - assert "Thought: t1" in combined - assert "Final Answer: done" in combined + assert result == ["system", "historic", "query", *runner._current_thoughts] diff --git a/api/tests/unit_tests/core/agent/test_fc_agent_runner.py b/api/tests/unit_tests/core/agent/test_fc_agent_runner.py index 89eb6c4674d67d..f7921806cdbe2c 100644 --- a/api/tests/unit_tests/core/agent/test_fc_agent_runner.py +++ b/api/tests/unit_tests/core/agent/test_fc_agent_runner.py @@ -1,9 +1,7 @@ import json from collections.abc import Iterator -from dataclasses import dataclass from datetime import UTC, datetime from decimal import Decimal -from typing import Any from unittest.mock import MagicMock import pytest @@ -26,7 +24,6 @@ from graphon.model_runtime.entities.llm_entities import LLMResult, LLMResultChunk, LLMResultChunkDelta, LLMUsage from graphon.model_runtime.entities.message_entities import ( AssistantPromptMessage, - DocumentPromptMessageContent, ImagePromptMessageContent, PromptMessageContentType, SystemPromptMessage, @@ -83,34 +80,12 @@ def _make_message(*, message_id: str = "m1", conversation_id: str = "conv1") -> ) -class DummyMessage: - def __init__(self, content: str | None = None, tool_calls: list[Any] | None = None): - self.content: str | None = content - self.tool_calls: list[Any] = tool_calls or [] +def _chunk(message, usage=None): + return LLMResultChunk(model="m", delta=LLMResultChunkDelta(index=0, message=message, usage=usage)) -@dataclass(eq=False, repr=False) -class DummyDelta: - message: DummyMessage | None = None - usage: LLMUsage | None = None - - -class DummyChunk: - def __init__(self, message: DummyMessage | None = None, usage: LLMUsage | None = None): - self.delta: DummyDelta = DummyDelta(message=message, usage=usage) - - -class DummyResult: - def __init__( - self, - message: DummyMessage | None = None, - usage: LLMUsage | None = None, - prompt_messages: list[DummyMessage] | None = None, - ): - self.message: DummyMessage | None = message - self.usage: LLMUsage | None = usage - self.prompt_messages: list[DummyMessage] = prompt_messages or [] - self.system_fingerprint: str = "" +def _result(message, usage=None): + return LLMResult(model="m", message=message, prompt_messages=[], usage=usage or LLMUsage.empty_usage()) # ============================== @@ -166,10 +141,6 @@ def runner(mocker: MockerFixture, sqlite_engine: Engine) -> Iterator[FunctionCal return_value=None, ) - # Patch streaming chunk models to avoid validation on dummy message objects - mocker.patch("core.agent.fc_agent_runner.LLMResultChunk", MagicMock) - mocker.patch("core.agent.fc_agent_runner.LLMResultChunkDelta", MagicMock) - app_config = MagicMock() app_config.app_id = "app" app_config.agent = MagicMock(max_iteration=2) @@ -361,35 +332,15 @@ def test_with_files_uses_image_detail_config(self, runner: FunctionCallAgentRunn # ============================== -class TestClearUserPromptImageMessages: - def test_clear_text_and_image_content(self, real_runner: FunctionCallAgentRunner) -> None: - text = TextPromptMessageContent(data="hello") - image = ImagePromptMessageContent(format="url", mime_type="image/png", url="https://example.test/image.png") - user_msg = UserPromptMessage(content=[text, image]) - - result = real_runner._clear_user_prompt_image_messages([user_msg]) - assert result == [UserPromptMessage(content="hello\n[image]")] - assert user_msg.content == [text, image] - - def test_clear_includes_file_placeholder(self, runner: FunctionCallAgentRunner): - text = TextPromptMessageContent(data="hello") - image = ImagePromptMessageContent(format="url", mime_type="image/png") - document = DocumentPromptMessageContent(format="url", mime_type="application/pdf") - - user_msg = UserPromptMessage(content=[text, image, document]) - - result = runner._clear_user_prompt_image_messages([user_msg]) - - assert result[0].content == "hello\n[image]\n[file]" - - def test_keeps_knowledge_retrieval_image_message(self, runner: FunctionCallAgentRunner): - text = TextPromptMessageContent(data="query") - image = ImagePromptMessageContent(format="url", mime_type="image/png") - user_msg = UserPromptMessage(name="knowledge_retrieval", content=[image, text]) - - result = runner._clear_user_prompt_image_messages([user_msg]) - - assert result[0].content == [image, text] +def test_prompt_keeps_images_after_tool_calls(real_runner, mocker): + image = ImagePromptMessageContent(format="png", mime_type="image/png", url="https://example.test/image.png") + history = [UserPromptMessage(content=[image, TextPromptMessageContent(data="describe")])] + real_runner.history_prompt_messages = history + real_runner._current_thoughts = [AssistantPromptMessage(content="checking", opaque_body=False)] + mocker.patch("core.agent.fc_agent_runner.AgentHistoryPromptTransform.get_prompt", return_value=history) + prompts = real_runner._organize_prompt_messages() + assert prompts[0] == history[0] + assert prompts[-1] == real_runner._current_thoughts[0] # ============================== @@ -505,8 +456,8 @@ def test_builds_image_contents_from_dataset_tool_preview_links( class TestRunMethod: def test_run_non_streaming_no_tool_calls(self, runner: FunctionCallAgentRunner): message = _make_message() - dummy_message = DummyMessage(content="hello") - result = DummyResult(message=dummy_message, usage=build_usage()) + dummy_message = AssistantPromptMessage(content="hello") + result = _result(message=dummy_message, usage=build_usage()) runner.model_instance.invoke_llm.return_value = result @@ -544,7 +495,7 @@ def record_detach(_session: Session, instance: object) -> None: event.listen(session, "persistent_to_detached", record_detach) content = [TextPromptMessageContent(data="hi")] - chunk = DummyChunk(message=DummyMessage(content=content), usage=build_usage()) + chunk = _chunk(message=AssistantPromptMessage(content=content), usage=build_usage()) def generator(): events.append("first-chunk") @@ -560,19 +511,16 @@ def test_run_streaming_tool_calls_list_content(self, runner: FunctionCallAgentRu message = _make_message() runner.stream_tool_call = True - tool_call = MagicMock() - tool_call.id = "1" - tool_call.function.name = "tool" - tool_call.function.arguments = json.dumps({"a": 1}) + tool_call = _tool_call("1", "tool", json.dumps({"a": 1})) content = [TextPromptMessageContent(data="hi")] - chunk = DummyChunk(message=DummyMessage(content=content, tool_calls=[tool_call]), usage=build_usage()) + chunk = _chunk(message=AssistantPromptMessage(content=content, tool_calls=[tool_call]), usage=build_usage()) def generator(): yield chunk - final_message = DummyMessage(content="done", tool_calls=[]) - final_result = DummyResult(message=final_message, usage=build_usage()) + final_message = AssistantPromptMessage(content="done", tool_calls=[]) + final_result = _result(message=final_message, usage=build_usage()) runner.model_instance.invoke_llm.side_effect = [generator(), final_result] @@ -582,8 +530,8 @@ def generator(): def test_run_non_streaming_list_content(self, runner: FunctionCallAgentRunner): message = _make_message() content = [TextPromptMessageContent(data="hi")] - dummy_message = DummyMessage(content=content) - result = DummyResult(message=dummy_message, usage=build_usage()) + dummy_message = AssistantPromptMessage(content=content) + result = _result(message=dummy_message, usage=build_usage()) runner.model_instance.invoke_llm.return_value = result @@ -591,46 +539,15 @@ def test_run_non_streaming_list_content(self, runner: FunctionCallAgentRunner): assert len(outputs) == 1 assert runner.save_agent_thought.call_args.kwargs["thought"] == "hi" - def test_run_streaming_tool_call_inputs_type_error(self, runner: FunctionCallAgentRunner, mocker: MockerFixture): - message = _make_message() - runner.stream_tool_call = True - - tool_call = MagicMock() - tool_call.id = "1" - tool_call.function.name = "tool" - tool_call.function.arguments = json.dumps({"a": 1}) - - chunk = DummyChunk(message=DummyMessage(content="hi", tool_calls=[tool_call]), usage=build_usage()) - - def generator(): - yield chunk - - runner.model_instance.invoke_llm.return_value = generator() - - real_dumps = json.dumps - - def flaky_dumps(obj, *args, **kwargs): - if kwargs.get("ensure_ascii") is False: - return real_dumps(obj, *args, **kwargs) - raise TypeError("boom") - - mocker.patch("core.agent.fc_agent_runner.json.dumps", side_effect=flaky_dumps) - - outputs = list(runner.run(runner.session, message, "query")) - assert len(outputs) == 1 - def test_run_with_missing_tool_instance(self, runner: FunctionCallAgentRunner): message = _make_message() - tool_call = MagicMock() - tool_call.id = "1" - tool_call.function.name = "missing" - tool_call.function.arguments = json.dumps({}) + tool_call = _tool_call("1", "missing", json.dumps({})) - dummy_message = DummyMessage(content="", tool_calls=[tool_call]) - result = DummyResult(message=dummy_message, usage=build_usage()) - final_message = DummyMessage(content="done", tool_calls=[]) - final_result = DummyResult(message=final_message, usage=build_usage()) + dummy_message = AssistantPromptMessage(content="", tool_calls=[tool_call]) + result = _result(message=dummy_message, usage=build_usage()) + final_message = AssistantPromptMessage(content="done", tool_calls=[]) + final_result = _result(message=final_message, usage=build_usage()) runner.model_instance.invoke_llm.side_effect = [result, final_result] @@ -640,14 +557,11 @@ def test_run_with_missing_tool_instance(self, runner: FunctionCallAgentRunner): def test_run_with_tool_instance_and_files(self, runner: FunctionCallAgentRunner, mocker: MockerFixture): message = _make_message() - tool_call = MagicMock() - tool_call.id = "1" - tool_call.function.name = "tool" - tool_call.function.arguments = json.dumps({"a": 1}) + tool_call = _tool_call("1", "tool", json.dumps({"a": 1})) - dummy_message = DummyMessage(content="", tool_calls=[tool_call]) - result = DummyResult(message=dummy_message, usage=build_usage()) - final_result = DummyResult(message=DummyMessage(content="done", tool_calls=[]), usage=build_usage()) + dummy_message = AssistantPromptMessage(content="", tool_calls=[tool_call]) + result = _result(message=dummy_message, usage=build_usage()) + final_result = _result(message=AssistantPromptMessage(content="done", tool_calls=[]), usage=build_usage()) runner.model_instance.invoke_llm.side_effect = [result, final_result] @@ -677,13 +591,10 @@ def test_run_max_iteration_error(self, runner: FunctionCallAgentRunner): message = _make_message() - tool_call = MagicMock() - tool_call.id = "1" - tool_call.function.name = "tool" - tool_call.function.arguments = "{}" + tool_call = _tool_call("1", "tool", "{}") - dummy_message = DummyMessage(content="", tool_calls=[tool_call]) - result = DummyResult(message=dummy_message, usage=build_usage()) + dummy_message = AssistantPromptMessage(content="", tool_calls=[tool_call]) + result = _result(message=dummy_message, usage=build_usage()) runner.model_instance.invoke_llm.return_value = result diff --git a/api/tests/unit_tests/core/agent/test_message_replay.py b/api/tests/unit_tests/core/agent/test_message_replay.py new file mode 100644 index 00000000000000..ffae756554ef01 --- /dev/null +++ b/api/tests/unit_tests/core/agent/test_message_replay.py @@ -0,0 +1,122 @@ +from copy import deepcopy + +import pytest +from sqlalchemy import select + +from core.agent.cot_chat_agent_runner import CotChatAgentRunner +from core.agent.fc_agent_runner import FunctionCallAgentRunner +from core.model_manager import ModelInstance +from graphon.model_runtime.entities import ( + AssistantPromptMessage, + LLMResult, + LLMResultChunk, + LLMResultChunkDelta, + LLMUsage, + TextPromptMessageContent, + ToolPromptMessage, +) +from models.model import MessageAgentThought +from services.agent.message_history import AgentMessageHistory +from services.workflow_run_agg import WorkflowRunAgg + + +@pytest.mark.usefixtures("agent_runtime_app") +@pytest.mark.parametrize("mode", ["blocking", "streaming", "react"]) +def test_tool_rounds_and_persisted_history_preserve_model_messages( + mode, sqlite_session, agent_records, agent_generate_entity, agent_queue_manager, monkeypatch +): + conversation, message = agent_records + app_config = agent_generate_entity.app_config + app_config.agent.max_iteration = 3 + runner_class = CotChatAgentRunner if mode == "react" else FunctionCallAgentRunner + runner = runner_class( + execution_driver=WorkflowRunAgg.run, + session=sqlite_session, + tenant_id=app_config.tenant_id, + application_generate_entity=agent_generate_entity, + conversation=conversation, + app_config=app_config, + model_config=agent_generate_entity.model_conf, + config=app_config.agent, + queue_manager=agent_queue_manager, + message=message, + user_id=agent_generate_entity.user_id, + model_instance=ModelInstance( + provider_model_bundle=agent_generate_entity.model_conf.provider_model_bundle, + model=agent_generate_entity.model_conf.model, + ), + ) + runner.stream_tool_call = mode == "streaming" + replies = [] + for index, state in enumerate(({"signature": "signed", "id": 9007199254740993}, False)): + text = 'Thought: check\nAction: {"action":"lookup","action_input":{}}' if mode == "react" else "check" + replies.append( + AssistantPromptMessage( + content=[ + TextPromptMessageContent(data=text, opaque_body={"block": index}), + TextPromptMessageContent(data="", opaque_body=0), + ], + opaque_body=state, + tool_calls=[] + if mode == "react" + else [ + AssistantPromptMessage.ToolCall( + id=f"original-{index}", + type="function", + function=AssistantPromptMessage.ToolCall.ToolCallFunction(name="lookup", arguments="{}"), + ) + ], + ) + ) + replies.append( + AssistantPromptMessage(content='{"action":"Final Answer","action_input":"done"}' if mode == "react" else "done") + ) + requests = [] + + def invoke(**kwargs): + requests.append(deepcopy(kwargs["prompt_messages"])) + reply = replies[len(requests) - 1].model_copy(deep=True) + if mode == "blocking": + return LLMResult(model="m", prompt_messages=[], message=reply, usage=LLMUsage.empty_usage()) + return (chunk for chunk in [LLMResultChunk(model="m", delta=LLMResultChunkDelta(index=0, message=reply))]) + + monkeypatch.setattr(runner.model_instance, "invoke_llm", invoke) + monkeypatch.setattr(runner, "_init_prompt_tools", lambda: ({}, [])) + list(runner.run(sqlite_session, message, "query", **({"inputs": {}} if mode == "react" else {}))) + assert len(requests) == 3 + assistants = [item for item in requests[-1] if isinstance(item, AssistantPromptMessage)] + assert assistants == replies[:2] + + rows = sqlite_session.scalars(select(MessageAgentThought).order_by(MessageAgentThought.position)).all() + restored = [item for row in rows for item in AgentMessageHistory.model_validate_json(row.model_messages).messages] + assert [item for item in restored if isinstance(item, AssistantPromptMessage)] == replies + if mode != "react": + assert [item.tool_call_id for item in restored if isinstance(item, ToolPromptMessage)] == [ + "original-0", + "original-1", + ] + runner.message = deepcopy(message) + runner.message.id = "next-message" + history = runner.organize_agent_history([], session=sqlite_session) + assert [item for item in history if isinstance(item, AssistantPromptMessage)] == replies + + +def test_history_keeps_empty_tool_results_and_completes_interrupted_calls(): + assistant = AssistantPromptMessage( + opaque_body=False, + tool_calls=[ + AssistantPromptMessage.ToolCall( + id="original", + type="function", + function=AssistantPromptMessage.ToolCall.ToolCallFunction(name="lookup", arguments="{}"), + ) + ], + ) + history = AgentMessageHistory(provider="p", model="m", messages=[assistant]) + restored = AgentMessageHistory.model_validate_json(history.model_dump_json()).restore(provider="p", model="m") + assert restored[0] == assistant + assert restored[1].tool_call_id == "original" + assert "interrupted" in restored[1].content + history.messages.append(ToolPromptMessage(tool_call_id="original", content="")) + assert history.restore(provider="p", model="m") == history.messages + assert history.restore(provider="other", model="m") is None diff --git a/api/tests/unit_tests/core/llm_generator/output_parser/test_structured_output.py b/api/tests/unit_tests/core/llm_generator/output_parser/test_structured_output.py index 548cc63000d905..f099131820acd7 100644 --- a/api/tests/unit_tests/core/llm_generator/output_parser/test_structured_output.py +++ b/api/tests/unit_tests/core/llm_generator/output_parser/test_structured_output.py @@ -388,29 +388,35 @@ def test_invoke_llm_with_structured_output_no_stream_prompt_based(self): assert result.structured_output == {"result": "success"} assert result.system_fingerprint == "fp_prompt" - def test_invoke_llm_with_structured_output_no_string_error(self): + def test_invoke_llm_with_structured_output_preserves_content_blocks(self): model_schema = make_model_schema("gpt-4") model_schema.features = [] model_schema.parameter_rules = [] llm_response = LLMResult( model=model_schema.model, - message=AssistantPromptMessage(content=[TextPromptMessageContent(data="not a string")]), + message=AssistantPromptMessage( + content=[ + TextPromptMessageContent(data='{"value":', opaque_body=False), + TextPromptMessageContent(data="1}"), + ], + opaque_body=0, + ), usage=LLMUsage.empty_usage(), ) model_instance = _ModelInstance(llm_response) - with pytest.raises(OutputParserError) as excinfo: - invoke_llm_with_structured_output( - provider="anthropic", - model_schema=model_schema, - model_instance=model_instance, - prompt_messages=[], - json_schema={}, - stream=False, - ) - assert "Failed to parse structured output, LLM result is not a string" in str(excinfo.value) + result = invoke_llm_with_structured_output( + provider="anthropic", + model_schema=model_schema, + model_instance=model_instance, + prompt_messages=[], + json_schema={}, + stream=False, + ) + assert result.structured_output == {"value": 1} + assert result.message == llm_response.message def test_invoke_llm_with_structured_output_stream(self): model_schema = make_model_schema("gpt-4") diff --git a/api/tests/unit_tests/core/prompt/test_agent_history_prompt_transform.py b/api/tests/unit_tests/core/prompt/test_agent_history_prompt_transform.py index 34aa7cf0248a4c..baa8a162b577b0 100644 --- a/api/tests/unit_tests/core/prompt/test_agent_history_prompt_transform.py +++ b/api/tests/unit_tests/core/prompt/test_agent_history_prompt_transform.py @@ -1,5 +1,6 @@ from unittest.mock import MagicMock +import pytest from pytest_mock import MockerFixture from core.prompt.agent_history_prompt_transform import AgentHistoryPromptTransform @@ -63,3 +64,30 @@ def side_effect_get_num_tokens(*args: object) -> int: result = transform.get_prompt() assert len(result) == 12 + + +@pytest.mark.parametrize("observation_name", ["agent_observation", "knowledge_retrieval"]) +def test_history_limit_keeps_complete_agent_turns(mocker: MockerFixture, observation_name: str) -> None: + history = [ + UserPromptMessage(content="Earlier query"), + AssistantPromptMessage(content="Earlier thought", opaque_body={"signature": "signed"}), + UserPromptMessage(content="Tool observation", name=observation_name), + AssistantPromptMessage(content="Earlier answer"), + UserPromptMessage(content="Latest query"), + AssistantPromptMessage(content="Latest answer"), + ] + model_config = make_model_config(provider="openai", model="gpt-4", mode="chat") + mocker.patch.object( + model_config.provider_model_bundle.model_type_instance, + "get_num_tokens", + side_effect=lambda _model, _credentials, messages: len(messages), + ) + transform = AgentHistoryPromptTransform( + model_config=model_config, + prompt_messages=[], + history_messages=history, + memory=make_token_buffer_memory(model_config), + ) + mocker.patch.object(transform, "_calculate_rest_token", return_value=4) + + assert transform.get_prompt() == history[-2:] diff --git a/api/tests/unit_tests/core/workflow/nodes/agent/test_message_history.py b/api/tests/unit_tests/core/workflow/nodes/agent/test_message_history.py new file mode 100644 index 00000000000000..f28e60ea65d4f8 --- /dev/null +++ b/api/tests/unit_tests/core/workflow/nodes/agent/test_message_history.py @@ -0,0 +1,142 @@ +import json + +import pytest + +from core.memory.token_buffer_memory import TokenBufferMemory +from core.model_manager import ModelInstance +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.workflow.nodes.agent.message_transformer import AgentMessageTransformer +from graphon.enums import BuiltinNodeTypes, WorkflowNodeExecutionStatus +from graphon.model_runtime.entities.message_entities import ( + AssistantPromptMessage, + TextPromptMessageContent, + ToolPromptMessage, + UserPromptMessage, +) +from graphon.node_events import StreamCompletedEvent +from models.enums import CreatorUserRole +from models.workflow import WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom +from services.agent.message_history import AgentMessageHistory +from services.agent.workflow_message_history import get_workflow_agent_history + + +@pytest.mark.usefixtures("agent_runtime_app") +@pytest.mark.parametrize("model", ["m", "different-model"]) +def test_workflow_history_round_trip_is_scoped_to_conversation_node_and_model( + model, sqlite_session, sqlite_engine, agent_records, agent_generate_entity +): + conversation, message = agent_records + message.workflow_run_id = "run" + message.answer = "visible answer" + message.answer_tokens = 1 + provider = agent_generate_entity.model_conf.provider + messages = [ + UserPromptMessage(content="original node query"), + AssistantPromptMessage( + content=[TextPromptMessageContent(data="", opaque_body={"id": 9007199254740993})], + opaque_body=False, + tool_calls=[ + AssistantPromptMessage.ToolCall( + id="original-call", + type="function", + function=AssistantPromptMessage.ToolCall.ToolCallFunction(name="lookup", arguments="{}"), + ) + ], + ), + ToolPromptMessage(tool_call_id="original-call", content=""), + AssistantPromptMessage(content="visible answer", opaque_body=0), + ] + history = AgentMessageHistory(provider=provider, model="m", messages=messages) + events = list( + AgentMessageTransformer().transform( + messages=( + item + for item in [ + ToolInvokeMessage( + type=ToolInvokeMessage.MessageType.JSON, + message=ToolInvokeMessage.JsonMessage( + json_object={ + "model_history": history.model_dump(mode="json"), + "public_value": 1, + } + ), + ) + ] + ), + tool_info={}, + parameters_for_log={}, + user_id="user", + tenant_id="tenant", + conversation_id=conversation.id, + node_type=BuiltinNodeTypes.AGENT, + node_id="agent", + node_execution_id="execution", + ) + ) + result = next(event.node_run_result for event in events if isinstance(event, StreamCompletedEvent)) + assert result.outputs["json"] == [{"public_value": 1}] + WorkflowNodeExecutionModel.__table__.create(sqlite_engine, checkfirst=True) + for index, (tenant, app, run, node) in enumerate( + [ + ("tenant", "app1", "run", "agent"), + ("other-tenant", "app1", "run", "agent"), + ("tenant", "other-app", "run", "agent"), + ("tenant", "app1", "other-run", "agent"), + ("tenant", "app1", "run", "other-node"), + ] + ): + sqlite_session.add( + WorkflowNodeExecutionModel( + id=f"execution-{index}", + tenant_id=tenant, + app_id=app, + workflow_id="workflow", + workflow_run_id=run, + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + index=index, + node_id=node, + node_type="agent", + title="Agent", + status=WorkflowNodeExecutionStatus.SUCCEEDED, + created_by_role=CreatorUserRole.ACCOUNT, + created_by="user", + process_data=json.dumps(result.process_data), + ) + ) + sqlite_session.commit() + memory = TokenBufferMemory( + conversation, + ModelInstance( + provider_model_bundle=agent_generate_entity.model_conf.provider_model_bundle, + model=model, + ), + ) + restored = get_workflow_agent_history( + memory=memory, + tenant_id="tenant", + app_id="app1", + node_id="agent", + provider=provider, + model=model, + message_limit=5, + ) + assert restored == ( + messages + if model == "m" + else [ + UserPromptMessage(content=message.query), + AssistantPromptMessage(content=message.answer), + ] + ) + assert ( + get_workflow_agent_history( + memory=memory, + tenant_id="other-tenant", + app_id="app1", + node_id="agent", + provider=provider, + model=model, + message_limit=5, + ) + == [] + ) diff --git a/api/uv.lock b/api/uv.lock index 36977a5b38bf9a..676a9b8747a41f 100644 --- a/api/uv.lock +++ b/api/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 3 +revision = 5 requires-python = "==3.13.*" resolution-markers = [ "sys_platform == 'win32'", @@ -1354,7 +1354,7 @@ dependencies = [ requires-dist = [ { name = "e2b", marker = "extra == 'server'", specifier = ">=2.55.1,<3.0.0" }, { name = "fastapi", marker = "extra == 'server'", specifier = "==0.136.0" }, - { name = "graphon", marker = "extra == 'server'", specifier = "==0.5.2" }, + { name = "graphon", marker = "extra == 'server'", git = "https://github.com/langgenius/graphon?rev=bfeca97726f91eab257e513357b72d456de851b2" }, { name = "grpcio", marker = "extra == 'server'", specifier = ">=1.60.0,<2.0.0" }, { name = "httpx", specifier = "==0.28.1" }, { name = "httpx2", specifier = ">=2.5.0,<3.0.0" }, @@ -1393,6 +1393,7 @@ docs = [ name = "dify-api" version = "1.17.1" source = { virtual = "." } +default-groups = ["kms", "storage", "tools", "trace-all", "vdb-all"] dependencies = [ { name = "aliyun-log-python-sdk" }, { name = "azure-identity" }, @@ -1705,7 +1706,7 @@ requires-dist = [ { name = "gmpy2", specifier = ">=2.3.0,<3.0.0" }, { name = "google-api-python-client", specifier = ">=2.201.0,<3.0.0" }, { name = "google-cloud-aiplatform", specifier = ">=2.3.0,<3.0.0" }, - { name = "graphon", git = "https://github.com/langgenius/graphon?rev=4fa22141774dff80e90e4dfa98ea3e6fbf0421bb" }, + { name = "graphon", git = "https://github.com/langgenius/graphon?rev=bfeca97726f91eab257e513357b72d456de851b2" }, { name = "gunicorn", specifier = ">=26.2.0,<27.0.0" }, { name = "httpx", extras = ["socks"], specifier = "==0.28.1" }, { name = "httpx-sse", specifier = "==0.4.3" }, @@ -3084,7 +3085,7 @@ httpx = [ [[package]] name = "graphon" version = "0.8.0" -source = { git = "https://github.com/langgenius/graphon?rev=4fa22141774dff80e90e4dfa98ea3e6fbf0421bb#4fa22141774dff80e90e4dfa98ea3e6fbf0421bb" } +source = { git = "https://github.com/langgenius/graphon?rev=bfeca97726f91eab257e513357b72d456de851b2#bfeca97726f91eab257e513357b72d456de851b2" } dependencies = [ { name = "charset-normalizer" }, { name = "defusedxml" }, diff --git a/dify-agent/docs/dify-agent/user-manual/history-layer/index.md b/dify-agent/docs/dify-agent/user-manual/history-layer/index.md index 1da5eea2088720..5c10b7be8e1d09 100644 --- a/dify-agent/docs/dify-agent/user-manual/history-layer/index.md +++ b/dify-agent/docs/dify-agent/user-manual/history-layer/index.md @@ -6,3 +6,9 @@ schemas. The runner explicitly loads initial message_history and writes the captured history on every terminal outcome, preserving compaction and partial responses. Current instructions are removed and unmatched trailing calls are marked interrupted on failure or cancellation. The reserved name is `history`. + +Dify model responses retain their original content blocks and provider state in +`ModelResponse.provider_details["dify_message"]` so tool rounds and restored runs +can replay signed or encrypted content without rebuilding it from display text. +Replay metadata is reused only when both the provider and model match the saved response. +Older history without this field continues to use its native message parts. diff --git a/dify-agent/pyproject.toml b/dify-agent/pyproject.toml index 11b6594604b717..6b0670f904060c 100644 --- a/dify-agent/pyproject.toml +++ b/dify-agent/pyproject.toml @@ -21,7 +21,7 @@ dify-agent-stub-server = "dify_agent.agent_stub.server.cli:main" server = [ "e2b>=2.55.1,<3.0.0", "fastapi==0.136.0", - "graphon==0.5.2", + "graphon==0.8.0", # grpcio is a runtime requirement of the openshell SDK that its wheel # does not declare; keep it pinned here until upstream fixes the metadata. "grpcio>=1.60.0,<2.0.0", @@ -35,6 +35,10 @@ server = [ "uvicorn[standard]==0.46.0", ] +[tool.uv.sources] +# shortcut: pin the joint replay candidate until Graphon releases incremental accumulation and container output fixes. +graphon = { git = "https://github.com/langgenius/graphon", rev = "bfeca97726f91eab257e513357b72d456de851b2" } + [tool.setuptools.packages.find] where = ["src"] include = ["dify_agent*", "shellctl*"] diff --git a/dify-agent/src/dify_agent/adapters/llm/model.py b/dify-agent/src/dify_agent/adapters/llm/model.py index 1f4c93a161ea8d..9026581d0d0180 100644 --- a/dify-agent/src/dify_agent/adapters/llm/model.py +++ b/dify-agent/src/dify_agent/adapters/llm/model.py @@ -32,6 +32,7 @@ UserPromptMessage, VideoPromptMessageContent, ) +from graphon.model_runtime.model_providers.base.large_language_model import LLMResultAccumulator from typing_extensions import assert_never, override from pydantic_ai._parts_manager import ModelResponsePartsManager @@ -159,6 +160,7 @@ async def request( ), response_model_name=self.model_name, provider_name_value=self.system, + model_provider=self.model_provider, ) async for _event in response: pass @@ -192,6 +194,7 @@ async def request_stream( ), response_model_name=self.model_name, provider_name_value=self.system, + model_provider=self.model_provider, ) yield response self._record_usage(response.dify_usage) @@ -212,7 +215,9 @@ def _build_request_input( model_request_parameters: ModelRequestParameters, ) -> _DifyRequestInput: return _DifyRequestInput( - prompt_messages=_map_messages_to_prompt_messages(messages, model_request_parameters), + prompt_messages=_map_messages_to_prompt_messages( + messages, model_request_parameters, model=self.model, model_provider=self.model_provider + ), model_parameters=_map_model_settings_to_parameters(model_settings), tools=_map_tool_definitions_to_prompt_tools(model_request_parameters), stop_sequences=_get_stop_sequences(model_settings), @@ -230,9 +235,11 @@ class DifyStreamedResponse(StreamedResponse): chunks: AsyncIterator[LLMResultChunk] response_model_name: str provider_name_value: str + model_provider: str _timestamp: datetime = field(default_factory=lambda: datetime.now(timezone.utc)) _embedded_thinking_parser: "_EmbeddedThinkingParser" = field(default_factory=lambda: _EmbeddedThinkingParser()) _dify_usage: LLMUsage | None = field(default=None, init=False, repr=False) + _result: LLMResultAccumulator = field(default_factory=LLMResultAccumulator, init=False, repr=False) @override async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]: @@ -244,19 +251,32 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]: if chunk.delta.finish_reason is not None: self.finish_reason: FinishReason | None = _normalize_finish_reason(chunk.delta.finish_reason) - for event in _chunk_to_stream_events( + events = _chunk_to_stream_events( self._parts_manager, chunk, self.provider_name_value, self._embedded_thinking_parser, chunk_sequence, - ): + ) + self._result.consume(chunk) + for event in events: yield event chunk_sequence += 1 for event in self._embedded_thinking_parser.flush(self._parts_manager, self.provider_name_value): yield event + @override + def get(self) -> ModelResponse: + response = super().get() + message = self._result.to_result(model=self.model_name, prompt_messages=[]).message + response.provider_details = { + **(response.provider_details or {}), + "dify_provider": self.model_provider, + "dify_message": message.model_dump(mode="json", exclude={"tool_calls"}), + } + return response + @property def dify_usage(self) -> LLMUsage | None: """Return Dify's complete usage for this model request.""" @@ -286,6 +306,9 @@ def timestamp(self) -> datetime: def _map_messages_to_prompt_messages( messages: Sequence[ModelMessage], model_request_parameters: ModelRequestParameters, + *, + model: str, + model_provider: str, ) -> list[PromptMessage]: prompt_messages: list[PromptMessage] = [] @@ -293,7 +316,13 @@ def _map_messages_to_prompt_messages( if isinstance(message, ModelRequest): prompt_messages.extend(_map_model_request_to_prompt_messages(message)) elif isinstance(message, ModelResponse): - assistant_message = _map_model_response_to_prompt_message(message) + assistant_message = _map_model_response_to_prompt_message( + message, + preserve_provider_message=( + message.model_name == model + and (message.provider_details or {}).get("dify_provider") == model_provider + ), + ) if assistant_message is not None: prompt_messages.append(assistant_message) else: @@ -390,14 +419,10 @@ def _map_tool_return_part_to_prompt_message(part: ToolReturnPart) -> ToolPromptM def _map_model_response_to_prompt_message( message: ModelResponse, + *, + preserve_provider_message: bool, ) -> AssistantPromptMessage | None: - """Map prior assistant output into daemon prompt history. - - The plugin daemon requires ``PromptMessage.content`` to be present even when - an assistant turn contains only tool calls. Tool-call-only assistant history - therefore uses the empty string instead of ``null`` so the second request in - a tool round trip remains schema-compatible. - """ + """Restore original provider content and state, using native parts for tool IDs.""" content_parts: list[PromptMessageContentUnionTypes] = [] tool_calls: list[AssistantPromptMessage.ToolCall] = [] @@ -424,6 +449,12 @@ def _map_model_response_to_prompt_message( else: raise UnexpectedModelBehavior(f"Unsupported response part for daemon adapter: {type(part).__name__}") + stored = (message.provider_details or {}).get("dify_message") + if preserve_provider_message and stored is not None: + restored = AssistantPromptMessage.model_validate(stored) + restored.tool_calls = tool_calls + return None if restored.is_empty() else restored + content = _normalize_prompt_content(content_parts) if content is None and not tool_calls: return None diff --git a/dify-agent/tests/local/dify_agent/adapters/llm/test_model.py b/dify-agent/tests/local/dify_agent/adapters/llm/test_model.py index eb6df3820544a0..aae10b6791dad7 100644 --- a/dify-agent/tests/local/dify_agent/adapters/llm/test_model.py +++ b/dify-agent/tests/local/dify_agent/adapters/llm/test_model.py @@ -29,6 +29,7 @@ from dify_agent.adapters.llm import DifyApiLLMProvider, DifyLLMAdapterModel from dify_agent.layers.execution_context import DifyExecutionContextLayerConfig +from dify_agent.layers.history import State as HistoryState from ._test_support import ( AssistantPromptMessage, @@ -561,6 +562,99 @@ def handler(_request: httpx.Request) -> httpx.Response: self.assertEqual(response.usage.input_tokens, 11) self.assertEqual(response.usage.output_tokens, 7) + async def test_restored_history_preserves_provider_message(self) -> None: + for streaming in (False, True): + for snapshot in ({"signature": "signed", "id": 9007199254740993}, {}, [], "", 0, False): + with self.subTest(streaming=streaming, snapshot=snapshot): + tool_call = AssistantPromptMessage.ToolCall( + id="call-1", + type="function", + function=AssistantPromptMessage.ToolCall.ToolCallFunction( + name="lookup", + arguments='{"id":9007199254740993}', + ), + ) + block = TextPromptMessageContent(data="", opaque_body={"signature": "block"}) + chunks = [ + LLMResultChunk( + model="demo-model", + delta=LLMResultChunkDelta( + index=0, + message=AssistantPromptMessage(content="before", opaque_body={"old": True}), + ), + ), + LLMResultChunk( + model="demo-model", + delta=LLMResultChunkDelta( + index=1, + message=AssistantPromptMessage( + content=[block], + tool_calls=[tool_call], + opaque_body=snapshot, + ), + ), + ), + LLMResultChunk( + model="demo-model", + delta=LLMResultChunkDelta( + index=2, + message=AssistantPromptMessage(content="after"), + ), + ), + ] + requests = [] + + def handler(request: httpx.Request) -> httpx.Response: + requests.append(json.loads(request.content)) + return build_stream_response(*(chunks if len(requests) == 1 else single_text_chunk("done"))) + + messages = [ModelRequest(parts=[UserPromptPart("look up")])] + async with self.mock_gateway_stream(httpx.MockTransport(handler)): + adapter = DifyLLMAdapterModel("demo-model", self.make_provider(), model_provider="openai") + if streaming: + async with adapter.request_stream(messages, None, ModelRequestParameters()) as stream: + _events = [event async for event in stream] + response = stream.get() + else: + response = await adapter.request(messages, None, ModelRequestParameters()) + + saved = HistoryState( + messages=[ + *messages, + response, + ModelRequest( + parts=[ + ToolReturnPart(tool_name="lookup", tool_call_id="call-1", content=""), + ] + ), + ] + ).model_dump_json() + restored = HistoryState.model_validate_json(saved).messages + fresh_adapter = DifyLLMAdapterModel("demo-model", self.make_provider(), model_provider="openai") + await fresh_adapter.request(restored, None, ModelRequestParameters()) + for next_provider, next_model in (("other-provider", "demo-model"), ("openai", "other-model")): + changed_adapter = DifyLLMAdapterModel( + next_model, self.make_provider(), model_provider=next_provider + ) + await changed_adapter.request(restored, None, ModelRequestParameters()) + + for request in requests[2:]: + changed = AssistantPromptMessage.model_validate(request["target"]["prompt_messages"][1]) + self.assertIsNone(changed.opaque_body) + self.assertEqual(changed.get_text_content(), "beforeafter") + replayed = AssistantPromptMessage.model_validate(requests[1]["target"]["prompt_messages"][1]) + self.assertEqual( + replayed.content, + [ + TextPromptMessageContent(data="before"), + block, + TextPromptMessageContent(data="after"), + ], + ) + self.assertEqual(replayed.tool_calls, [tool_call]) + self.assertEqual(replayed.opaque_body, snapshot) + self.assertIs(type(replayed.opaque_body), type(snapshot)) + async def test_request_stream_splits_embedded_thinking_tags_from_text_content_parts(self) -> None: def handler(_request: httpx.Request) -> httpx.Response: return build_stream_response( diff --git a/dify-agent/tests/local/test_packaging.py b/dify-agent/tests/local/test_packaging.py index a430fdff9ff134..2a2d3fc80f19b8 100644 --- a/dify-agent/tests/local/test_packaging.py +++ b/dify-agent/tests/local/test_packaging.py @@ -18,7 +18,7 @@ SERVER_RUNTIME_DEPENDENCIES = { "e2b>=2.55.1,<3.0.0", "fastapi==0.136.0", - "graphon==0.5.2", + "graphon==0.8.0", "grpcio>=1.60.0,<2.0.0", "jsonschema>=4.23.0,<5.0.0", "jwcrypto>=1.5.6,<2", diff --git a/dify-agent/uv.lock b/dify-agent/uv.lock index e415920441c7f9..10578de08f5eb9 100644 --- a/dify-agent/uv.lock +++ b/dify-agent/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 3 +revision = 5 requires-python = ">=3.13, <4.0" resolution-markers = [ "python_full_version >= '3.14' and sys_platform == 'win32'", @@ -580,7 +580,7 @@ docs = [ requires-dist = [ { name = "e2b", marker = "extra == 'server'", specifier = ">=2.55.1,<3.0.0" }, { name = "fastapi", marker = "extra == 'server'", specifier = "==0.136.0" }, - { name = "graphon", marker = "extra == 'server'", specifier = "==0.5.2" }, + { name = "graphon", marker = "extra == 'server'", git = "https://github.com/langgenius/graphon?rev=bfeca97726f91eab257e513357b72d456de851b2" }, { name = "grpcio", marker = "extra == 'server'", specifier = ">=1.60.0,<2.0.0" }, { name = "httpx", specifier = "==0.28.1" }, { name = "httpx2", specifier = ">=2.5.0,<3.0.0" }, @@ -813,31 +813,29 @@ wheels = [ [[package]] name = "graphon" -version = "0.5.2" -source = { registry = "https://pypi.org/simple" } +version = "0.8.0" +source = { git = "https://github.com/langgenius/graphon?rev=bfeca97726f91eab257e513357b72d456de851b2#bfeca97726f91eab257e513357b72d456de851b2" } dependencies = [ { name = "charset-normalizer" }, + { name = "defusedxml" }, { name = "httpx" }, { name = "json-repair" }, { name = "jsonschema" }, + { name = "odfdo" }, { name = "orjson" }, { name = "pandas", extra = ["excel"] }, { name = "pydantic" }, { name = "pydantic-extra-types" }, - { name = "pypandoc" }, + { name = "pypdf" }, { name = "pypdfium2" }, { name = "python-docx" }, { name = "pyyaml" }, { name = "tiktoken" }, { name = "transformers" }, { name = "typing-extensions" }, - { name = "unstructured", extra = ["docx", "epub", "md", "ppt", "pptx"] }, + { name = "unstructured", extra = ["docx", "md", "ppt", "pptx"] }, { name = "webvtt-py" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/c2/16/f183da187414c335be67f52f6a1b7c2a33bf0b1d5090eda7e6c92d42d94a/graphon-0.5.2.tar.gz", hash = "sha256:d66a9edcd883766bd50e94f84a691c92ce536ea60e721552089e83ac8e94bf68", size = 269773, upload-time = "2026-06-16T04:06:22.074Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/2f/e6/36a3981cd44e7a40a7cd7d374e26f01e02dd49410c5fbbd7df248750d5fb/graphon-0.5.2-py3-none-any.whl", hash = "sha256:11f89399e67ed1ddd2ce1c336accd9c4ad5b8fe2741f9167e6085af0b325cd14", size = 381908, upload-time = "2026-06-16T04:06:20.453Z" }, -] [[package]] name = "griffelib" @@ -1717,6 +1715,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/58/78/548fb8e07b1a341746bfbecb32f2c268470f45fa028aacdbd10d9bc73aab/numpy-2.4.4-cp314-cp314t-win_arm64.whl", hash = "sha256:ba203255017337d39f89bdd58417f03c4426f12beed0440cfd933cb15f8669c7", size = 10566643, upload-time = "2026-03-29T13:21:34.339Z" }, ] +[[package]] +name = "odfdo" +version = "3.27.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "lxml" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/63/41/fdad16aa8962a6a6896c960755cabada1df89fa6be78ed0377e3921717fb/odfdo-3.27.2.tar.gz", hash = "sha256:8a55f373d03e18d177e8a1b5a3eb51b639bfdf9568586882f293bc52d206f37b", size = 342524, upload-time = "2026-10-17T11:43:55.953Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7d/b1/ec2cc49162afee4bd30969b25862187d99f191cdd6a48dc92aece4695ea0/odfdo-3.27.2-py3-none-any.whl", hash = "sha256:9a092b10b61562f396df9ff3b20db83f51c85c36409e865151d3ce231f543dc9", size = 440596, upload-time = "2026-10-17T11:43:54.479Z" }, +] + [[package]] name = "odfpy" version = "1.4.1" @@ -2478,28 +2489,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d6/54/da572c98c0b77626a91b5d3b89f0231d8bff5125c225420908632f8b342d/pymdown_extensions-11.0.1-py3-none-any.whl", hash = "sha256:db3943a62bab7e03af1364f0c4083e64b91fb097675a4b6cceccfbe9a77e5eb2", size = 269455, upload-time = "2026-07-02T17:59:21.271Z" }, ] -[[package]] -name = "pypandoc" -version = "1.17" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/ea/d6/410615fc433e5d1eacc00db2044ae2a9c82302df0d35366fe2bd15de024d/pypandoc-1.17.tar.gz", hash = "sha256:51179abfd6e582a25ed03477541b48836b5bba5a4c3b282a547630793934d799", size = 69071, upload-time = "2026-03-14T22:39:07.21Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/0c/86/e2ffa604eacfbec3f430b1d850e7e04c4101eca1a5828f9ae54bf51dfba4/pypandoc-1.17-py3-none-any.whl", hash = "sha256:01fdbffa61edb9f8e82e8faad6954efcb7b6f8f0634aead4d89e322a00225a67", size = 23554, upload-time = "2026-03-14T22:38:46.007Z" }, -] - -[[package]] -name = "pypandoc-binary" -version = "1.17" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/80/85/681a54111f0948821a5cf87ce30a88bb0a3f6848af5112c912abac4a2b77/pypandoc_binary-1.17-py3-none-macosx_10_9_x86_64.whl", hash = "sha256:734726dc618ef276343e272e1a6b4567e59c2ef9ef41d5533042deac3b0531f1", size = 25553945, upload-time = "2026-03-14T22:38:47.91Z" }, - { url = "https://files.pythonhosted.org/packages/15/58/8fd107c68522957868c1e785fbea7595608df118e440e424d189668294df/pypandoc_binary-1.17-py3-none-macosx_11_0_arm64.whl", hash = "sha256:fcfd28f347ed998dda28823fc6bc24f9310e7fdf3ddceaf925bf0563a100ab5b", size = 25553944, upload-time = "2026-03-14T22:38:50.74Z" }, - { url = "https://files.pythonhosted.org/packages/f4/27/ac1078239aae14b94c51975b7f46ad8e099e47d7ae26c175a5486b1c0099/pypandoc_binary-1.17-py3-none-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:d6b620b21c9374e3e48aabd518492bf0776b148442ee28816f6aaf52da3d4387", size = 34460960, upload-time = "2026-03-14T22:38:53.391Z" }, - { url = "https://files.pythonhosted.org/packages/8d/7f/1e5612b52900ebe590862dabeadf546f739b27527dcd8bfd632f8adac1be/pypandoc_binary-1.17-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9ada156cb980cd54fd6534231788e668c00dbb591cbd24f0be0bd86812eb8788", size = 36867598, upload-time = "2026-03-14T22:38:56.351Z" }, - { url = "https://files.pythonhosted.org/packages/3b/31/a5a867159c4080e5d368f4a53540a727501a2f31affc297dc8e0fced96a7/pypandoc_binary-1.17-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:2f439dcd211183bb3460253ca4511101df6e1acf4a01f45f5617e1fa2ad24279", size = 36867584, upload-time = "2026-03-14T22:38:59.087Z" }, - { url = "https://files.pythonhosted.org/packages/0d/2d/6a51cd4e54bdf132c19416801077c34bd40ba182e85d843360d36ae03a2d/pypandoc_binary-1.17-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:f6e6d3e4cfafbe23189a08db3d41f8def260bacd6e7e382bceadab7ba1f17da6", size = 34460949, upload-time = "2026-03-14T22:39:01.71Z" }, -] - [[package]] name = "pypdf" version = "6.19.0" @@ -3594,9 +3583,6 @@ wheels = [ docx = [ { name = "python-docx" }, ] -epub = [ - { name = "pypandoc-binary", marker = "sys_platform != 'win32'" }, -] md = [ { name = "markdown" }, ] From 3b7bc8309487c1230fe59867e3b4b0036d7cde06 Mon Sep 17 00:00:00 2001 From: WH-2099 Date: Wed, 28 Oct 2026 05:34:01 +0000 Subject: [PATCH 2/4] fix(agent): bound replay history and validate complete tool calls --- api/core/memory/token_buffer_memory.py | 16 +- api/pyproject.toml | 2 +- .../agent/workflow_message_history.py | 9 +- .../nodes/agent/test_message_history.py | 146 +++++++++++++++++- api/uv.lock | 6 +- .../user-manual/history-layer/index.md | 2 + dify-agent/pyproject.toml | 2 +- .../src/dify_agent/adapters/llm/model.py | 23 +-- .../dify_agent/adapters/llm/test_model.py | 47 +++++- dify-agent/uv.lock | 4 +- 10 files changed, 217 insertions(+), 40 deletions(-) diff --git a/api/core/memory/token_buffer_memory.py b/api/core/memory/token_buffer_memory.py index cc7c0a6e06f957..745b168d6abf1d 100644 --- a/api/core/memory/token_buffer_memory.py +++ b/api/core/memory/token_buffer_memory.py @@ -83,11 +83,16 @@ def to_prompt_message(self) -> PromptMessage: class PreparedHistory: """Detached history that can restore attachments and count tokens without its read session. + Each turn belongs to one database Message, including its restored agent/tool messages. File factories retain ownership of their short metadata lookups. Rendering may perform storage/network I/O, so callers should first close the history session. """ - prompts: tuple[HistoryPrompt | PromptMessage, ...] + turns: tuple[tuple[HistoryPrompt | PromptMessage, ...], ...] + + @property + def prompts(self) -> tuple[HistoryPrompt | PromptMessage, ...]: + return tuple(prompt for turn in self.turns for prompt in turn) def to_prompt_messages(self) -> list[PromptMessage]: return [prompt.to_prompt_message() if isinstance(prompt, HistoryPrompt) else prompt for prompt in self.prompts] @@ -262,12 +267,13 @@ def load_history( ).all(): assistant_files_by_message[message_file.message_id].append(message_file) - prompts: list[HistoryPrompt | PromptMessage] = [] + turns: list[tuple[HistoryPrompt | PromptMessage, ...]] = [] for message in messages: if message_loader is not None: restored = message_loader(message) if restored is not None: - prompts.extend(restored) + if restored: + turns.append(tuple(restored)) continue user_files = user_files_by_message.get(message.id, []) assistant_files = assistant_files_by_message.get(message.id, []) @@ -291,6 +297,7 @@ def load_history( if user_files or assistant_files else None ) + prompts: list[HistoryPrompt | PromptMessage] = [] for files, text_content, is_user_message in ( (user_files, message.query, True), (assistant_files, message.answer, False), @@ -304,7 +311,8 @@ def load_history( file_config=file_config, ) ) - return PreparedHistory(prompts=tuple(prompts)) + turns.append(tuple(prompts)) + return PreparedHistory(turns=tuple(turns)) def get_history_prompt_messages( self, max_token_limit: int = 2000, message_limit: int | None = None diff --git a/api/pyproject.toml b/api/pyproject.toml index 33b61486db1cfb..346cde5ae595e3 100644 --- a/api/pyproject.toml +++ b/api/pyproject.toml @@ -64,7 +64,7 @@ exclude = ["providers/vdb/__pycache__", "providers/trace/__pycache__"] dify-agent = { path = "../dify-agent", editable = true } flask-restx = { git = "https://github.com/asukaminato0721/flask-restx", rev = "27758e26f8f740d7525d5039c51a9e524b6e2b68" } # shortcut: pin the joint replay candidate until Graphon releases incremental accumulation and container output fixes. -graphon = { git = "https://github.com/langgenius/graphon", rev = "bfeca97726f91eab257e513357b72d456de851b2" } +graphon = { git = "https://github.com/langgenius/graphon", rev = "86f1e11b199bb8ad9111ed49d42edea82a16b7fa" } dify-vdb-alibabacloud-mysql = { workspace = true } dify-vdb-analyticdb = { workspace = true } dify-vdb-baidu = { workspace = true } diff --git a/api/services/agent/workflow_message_history.py b/api/services/agent/workflow_message_history.py index 813b128c2f915d..34d0c0eb6c26d3 100644 --- a/api/services/agent/workflow_message_history.py +++ b/api/services/agent/workflow_message_history.py @@ -23,7 +23,9 @@ def get_workflow_agent_history( provider: str, model: str, message_limit: int | None, + max_token_limit: int = 2000, ) -> list[PromptMessage]: + """Trim oldest complete turns using the model's configured token counter.""" session_maker = sessionmaker(bind=db.engine, expire_on_commit=False) repository = DifyAPIRepositoryFactory.create_api_workflow_node_execution_repository(session_maker) @@ -60,4 +62,9 @@ def load_message(message: Message) -> list[PromptMessage] | None: message_limit=message_limit, message_loader=load_message, ) - return history.to_prompt_messages() + prompt_messages = history.to_prompt_messages() + for turn in history.turns: + if memory.model_instance.get_llm_num_tokens(prompt_messages) <= max_token_limit: + break + del prompt_messages[: len(turn)] + return prompt_messages diff --git a/api/tests/unit_tests/core/workflow/nodes/agent/test_message_history.py b/api/tests/unit_tests/core/workflow/nodes/agent/test_message_history.py index f28e60ea65d4f8..ec1f9f02f8279e 100644 --- a/api/tests/unit_tests/core/workflow/nodes/agent/test_message_history.py +++ b/api/tests/unit_tests/core/workflow/nodes/agent/test_message_history.py @@ -1,20 +1,29 @@ import json +from collections.abc import Callable, Sequence +from datetime import timedelta +from typing import Literal import pytest +from sqlalchemy.engine import Engine +from sqlalchemy.orm import Session +from core.app.entities.app_invoke_entities import AgentChatAppGenerateEntity from core.memory.token_buffer_memory import TokenBufferMemory from core.model_manager import ModelInstance +from core.plugin.impl.model import PluginModelClient from core.tools.entities.tool_entities import ToolInvokeMessage from core.workflow.nodes.agent.message_transformer import AgentMessageTransformer from graphon.enums import BuiltinNodeTypes, WorkflowNodeExecutionStatus from graphon.model_runtime.entities.message_entities import ( AssistantPromptMessage, + PromptMessage, TextPromptMessageContent, ToolPromptMessage, UserPromptMessage, ) from graphon.node_events import StreamCompletedEvent from models.enums import CreatorUserRole +from models.model import Conversation, Message from models.workflow import WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom from services.agent.message_history import AgentMessageHistory from services.agent.workflow_message_history import get_workflow_agent_history @@ -23,14 +32,18 @@ @pytest.mark.usefixtures("agent_runtime_app") @pytest.mark.parametrize("model", ["m", "different-model"]) def test_workflow_history_round_trip_is_scoped_to_conversation_node_and_model( - model, sqlite_session, sqlite_engine, agent_records, agent_generate_entity -): + model: str, + sqlite_session: Session, + sqlite_engine: Engine, + agent_records: tuple[Conversation, Message], + agent_generate_entity: AgentChatAppGenerateEntity, +) -> None: conversation, message = agent_records message.workflow_run_id = "run" message.answer = "visible answer" message.answer_tokens = 1 provider = agent_generate_entity.model_conf.provider - messages = [ + messages: list[PromptMessage] = [ UserPromptMessage(content="original node query"), AssistantPromptMessage( content=[TextPromptMessageContent(data="", opaque_body={"id": 9007199254740993})], @@ -75,7 +88,9 @@ def test_workflow_history_round_trip_is_scoped_to_conversation_node_and_model( ) result = next(event.node_run_result for event in events if isinstance(event, StreamCompletedEvent)) assert result.outputs["json"] == [{"public_value": 1}] - WorkflowNodeExecutionModel.__table__.create(sqlite_engine, checkfirst=True) + WorkflowNodeExecutionModel.metadata.tables[WorkflowNodeExecutionModel.__tablename__].create( + sqlite_engine, checkfirst=True + ) for index, (tenant, app, run, node) in enumerate( [ ("tenant", "app1", "run", "agent"), @@ -140,3 +155,126 @@ def test_workflow_history_round_trip_is_scoped_to_conversation_node_and_model( ) == [] ) + + +@pytest.mark.usefixtures("agent_runtime_app") +@pytest.mark.parametrize("strategy", ["function_calling", "react"]) +@pytest.mark.parametrize( + ("max_token_limit", "plugin_counting", "retained"), + [ + pytest.param(None, True, "latest", id="default-budget"), + pytest.param(1599, True, "none", id="oversized-latest-turn"), + pytest.param(None, False, "all", id="disabled-token-counting"), + ], +) +def test_workflow_history_budget_preserves_database_turns( + strategy: Literal["function_calling", "react"], + max_token_limit: int | None, + plugin_counting: bool, + retained: Literal["latest", "none", "all"], + sqlite_session: Session, + sqlite_engine: Engine, + agent_records: tuple[Conversation, Message], + agent_generate_entity: AgentChatAppGenerateEntity, + config_overrides: Callable[..., None], + monkeypatch: pytest.MonkeyPatch, +) -> None: + config_overrides(PLUGIN_BASED_TOKEN_COUNTING_ENABLED=plugin_counting) + conversation, older = agent_records + older.answer = "older answer" + older.answer_tokens = 1 + latest = Message( + id="latest-message", + app_id=older.app_id, + conversation_id=conversation.id, + parent_message_id=older.id, + workflow_run_id="latest-run", + inputs={}, + query="latest query", + message={}, + message_unit_price=0, + answer="latest answer", + answer_tokens=1, + answer_unit_price=0, + currency="USD", + from_source=older.from_source, + created_at=older.created_at + timedelta(seconds=1), + ) + provider = agent_generate_entity.model_conf.provider + tool_call = AssistantPromptMessage.ToolCall( + id="lookup-call", + type="function", + function=AssistantPromptMessage.ToolCall.ToolCallFunction(name="lookup", arguments="{}"), + ) + messages: list[PromptMessage] = [ + UserPromptMessage(content=latest.query), + AssistantPromptMessage( + content=[TextPromptMessageContent(data="thought", opaque_body={"id": 9007199254740993})], + opaque_body=False, + tool_calls=[tool_call] if strategy == "function_calling" else [], + ), + ToolPromptMessage(tool_call_id=tool_call.id, content="") + if strategy == "function_calling" + else UserPromptMessage(content="Observation: lookup result"), + AssistantPromptMessage(content=latest.answer, opaque_body=0), + ] + history = AgentMessageHistory(provider=provider, model="m", messages=messages) + WorkflowNodeExecutionModel.metadata.tables[WorkflowNodeExecutionModel.__tablename__].create( + sqlite_engine, checkfirst=True + ) + sqlite_session.add_all( + [ + latest, + WorkflowNodeExecutionModel( + id="latest-execution", + tenant_id="tenant", + app_id=older.app_id, + workflow_id="workflow", + workflow_run_id=latest.workflow_run_id, + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + index=0, + node_id="agent", + node_type="agent", + title="Agent", + status=WorkflowNodeExecutionStatus.SUCCEEDED, + created_by_role=CreatorUserRole.ACCOUNT, + created_by="user", + process_data=json.dumps({"model_history": history.model_dump(mode="json")}), + ), + ] + ) + sqlite_session.commit() + memory = TokenBufferMemory( + conversation, + ModelInstance(provider_model_bundle=agent_generate_entity.model_conf.provider_model_bundle, model="m"), + ) + if plugin_counting: + + def count_tokens(prompt_messages: Sequence[PromptMessage]) -> int: + return 400 * len(prompt_messages) + + monkeypatch.setattr(memory.model_instance, "get_llm_num_tokens", count_tokens) + else: + + def reject_token_request(*_args: object, **_kwargs: object) -> None: + pytest.fail("Disabled token counting must not invoke the provider") + + monkeypatch.setattr(PluginModelClient, "get_llm_num_tokens", reject_token_request) + + budget = {} if max_token_limit is None else {"max_token_limit": max_token_limit} + restored = get_workflow_agent_history( + memory=memory, + tenant_id="tenant", + app_id=older.app_id, + node_id="agent", + provider=provider, + model="m", + message_limit=5, + **budget, + ) + expected: dict[str, list[PromptMessage]] = { + "latest": messages, + "none": [], + "all": [UserPromptMessage(content=older.query), AssistantPromptMessage(content=older.answer), *messages], + } + assert restored == expected[retained] diff --git a/api/uv.lock b/api/uv.lock index 676a9b8747a41f..07f62d5dcbd593 100644 --- a/api/uv.lock +++ b/api/uv.lock @@ -1354,7 +1354,7 @@ dependencies = [ requires-dist = [ { name = "e2b", marker = "extra == 'server'", specifier = ">=2.55.1,<3.0.0" }, { name = "fastapi", marker = "extra == 'server'", specifier = "==0.136.0" }, - { name = "graphon", marker = "extra == 'server'", git = "https://github.com/langgenius/graphon?rev=bfeca97726f91eab257e513357b72d456de851b2" }, + { name = "graphon", marker = "extra == 'server'", git = "https://github.com/langgenius/graphon?rev=86f1e11b199bb8ad9111ed49d42edea82a16b7fa" }, { name = "grpcio", marker = "extra == 'server'", specifier = ">=1.60.0,<2.0.0" }, { name = "httpx", specifier = "==0.28.1" }, { name = "httpx2", specifier = ">=2.5.0,<3.0.0" }, @@ -1706,7 +1706,7 @@ requires-dist = [ { name = "gmpy2", specifier = ">=2.3.0,<3.0.0" }, { name = "google-api-python-client", specifier = ">=2.201.0,<3.0.0" }, { name = "google-cloud-aiplatform", specifier = ">=2.3.0,<3.0.0" }, - { name = "graphon", git = "https://github.com/langgenius/graphon?rev=bfeca97726f91eab257e513357b72d456de851b2" }, + { name = "graphon", git = "https://github.com/langgenius/graphon?rev=86f1e11b199bb8ad9111ed49d42edea82a16b7fa" }, { name = "gunicorn", specifier = ">=26.2.0,<27.0.0" }, { name = "httpx", extras = ["socks"], specifier = "==0.28.1" }, { name = "httpx-sse", specifier = "==0.4.3" }, @@ -3085,7 +3085,7 @@ httpx = [ [[package]] name = "graphon" version = "0.8.0" -source = { git = "https://github.com/langgenius/graphon?rev=bfeca97726f91eab257e513357b72d456de851b2#bfeca97726f91eab257e513357b72d456de851b2" } +source = { git = "https://github.com/langgenius/graphon?rev=86f1e11b199bb8ad9111ed49d42edea82a16b7fa#86f1e11b199bb8ad9111ed49d42edea82a16b7fa" } dependencies = [ { name = "charset-normalizer" }, { name = "defusedxml" }, diff --git a/dify-agent/docs/dify-agent/user-manual/history-layer/index.md b/dify-agent/docs/dify-agent/user-manual/history-layer/index.md index 5c10b7be8e1d09..565ba69419c8eb 100644 --- a/dify-agent/docs/dify-agent/user-manual/history-layer/index.md +++ b/dify-agent/docs/dify-agent/user-manual/history-layer/index.md @@ -12,3 +12,5 @@ Dify model responses retain their original content blocks and provider state in can replay signed or encrypted content without rebuilding it from display text. Replay metadata is reused only when both the provider and model match the saved response. Older history without this field continues to use its native message parts. +Model streams carry complete tool calls once each, with nonempty IDs unique within the response. +The Dify adapter preserves those IDs and rejects invalid batches before creating native tool events. diff --git a/dify-agent/pyproject.toml b/dify-agent/pyproject.toml index 6b0670f904060c..de47f8f66f6d7b 100644 --- a/dify-agent/pyproject.toml +++ b/dify-agent/pyproject.toml @@ -37,7 +37,7 @@ server = [ [tool.uv.sources] # shortcut: pin the joint replay candidate until Graphon releases incremental accumulation and container output fixes. -graphon = { git = "https://github.com/langgenius/graphon", rev = "bfeca97726f91eab257e513357b72d456de851b2" } +graphon = { git = "https://github.com/langgenius/graphon", rev = "86f1e11b199bb8ad9111ed49d42edea82a16b7fa" } [tool.setuptools.packages.find] where = ["src"] diff --git a/dify-agent/src/dify_agent/adapters/llm/model.py b/dify-agent/src/dify_agent/adapters/llm/model.py index 9026581d0d0180..70ee32c57e4f38 100644 --- a/dify-agent/src/dify_agent/adapters/llm/model.py +++ b/dify-agent/src/dify_agent/adapters/llm/model.py @@ -243,8 +243,8 @@ class DifyStreamedResponse(StreamedResponse): @override async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]: - chunk_sequence = 0 async for chunk in self.chunks: + self._result.consume(chunk) if chunk.delta.usage is not None: self._dify_usage = chunk.delta.usage self._usage: RequestUsage = _map_usage(chunk.delta.usage) @@ -256,12 +256,9 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]: chunk, self.provider_name_value, self._embedded_thinking_parser, - chunk_sequence, ) - self._result.consume(chunk) for event in events: yield event - chunk_sequence += 1 for event in self._embedded_thinking_parser.flush(self._parts_manager, self.provider_name_value): yield event @@ -649,21 +646,11 @@ def _normalize_finish_reason(finish_reason: str) -> FinishReason: return "error" -def _normalize_tool_call_id(tool_call_id: str | None) -> str | None: - if tool_call_id is None: - return None - normalized = tool_call_id.strip() - if not normalized or normalized.lower() in {"none", "null"}: - return None - return normalized - - def _chunk_to_stream_events( parts_manager: ModelResponsePartsManager, chunk: LLMResultChunk, provider_name: str, embedded_thinking_parser: "_EmbeddedThinkingParser", - chunk_sequence: int, ) -> list[ModelResponseStreamEvent]: events: list[ModelResponseStreamEvent] = [] message = chunk.delta.message @@ -678,15 +665,13 @@ def _chunk_to_stream_events( else: events.append(parts_manager.handle_part(vendor_part_id=None, part=part)) - for index, tool_call in enumerate(message.tool_calls): - tool_call_id = _normalize_tool_call_id(tool_call.id) - vendor_id = tool_call_id or f"chunk-{chunk_sequence}-tool-{index}" + for tool_call in message.tool_calls: events.append( parts_manager.handle_tool_call_part( - vendor_part_id=vendor_id, + vendor_part_id=tool_call.id, tool_name=tool_call.function.name, args=tool_call.function.arguments, - tool_call_id=tool_call_id or vendor_id, + tool_call_id=tool_call.id, provider_name=provider_name, ) ) diff --git a/dify-agent/tests/local/dify_agent/adapters/llm/test_model.py b/dify-agent/tests/local/dify_agent/adapters/llm/test_model.py index aae10b6791dad7..8d28c783d9f5cd 100644 --- a/dify-agent/tests/local/dify_agent/adapters/llm/test_model.py +++ b/dify-agent/tests/local/dify_agent/adapters/llm/test_model.py @@ -776,7 +776,7 @@ def handler(_request: httpx.Request) -> httpx.Response: self.assertEqual(response.parts[2].part_kind, "text") self.assertEqual(cast(TextPart, response.parts[2]).content, "world") - async def test_request_stream_assigns_fallback_ids_to_tool_calls_without_ids(self) -> None: + async def test_request_stream_preserves_ids_of_complete_tool_calls(self) -> None: def handler(_request: httpx.Request) -> httpx.Response: return build_stream_response( LLMResultChunk( @@ -787,7 +787,7 @@ def handler(_request: httpx.Request) -> httpx.Response: content="", tool_calls=[ AssistantPromptMessage.ToolCall( - id=None, # pyright: ignore[reportArgumentType] + id="call-1", type="function", function=AssistantPromptMessage.ToolCall.ToolCallFunction( name="shell_run", @@ -806,7 +806,7 @@ def handler(_request: httpx.Request) -> httpx.Response: content="", tool_calls=[ AssistantPromptMessage.ToolCall( - id=None, # pyright: ignore[reportArgumentType] + id="call-2", type="function", function=AssistantPromptMessage.ToolCall.ToolCallFunction( name="shell_run", @@ -836,11 +836,48 @@ def handler(_request: httpx.Request) -> httpx.Response: self.assertTrue(events) self.assertEqual([part.part_kind for part in response.parts], ["tool-call", "tool-call"]) - self.assertEqual(cast(ToolCallPart, response.parts[0]).tool_call_id, "chunk-0-tool-0") - self.assertEqual(cast(ToolCallPart, response.parts[1]).tool_call_id, "chunk-1-tool-0") + self.assertEqual(cast(ToolCallPart, response.parts[0]).tool_call_id, "call-1") + self.assertEqual(cast(ToolCallPart, response.parts[1]).tool_call_id, "call-2") self.assertEqual(cast(ToolCallPart, response.parts[0]).args, '{"script":"lookup find"}') self.assertEqual(cast(ToolCallPart, response.parts[1]).args, '{"script":"lookup out"}') + async def test_request_stream_rejects_invalid_ids_before_creating_tool_parts(self) -> None: + for call_ids in ([""], ["same", "same"]): + with self.subTest(call_ids=call_ids): + + def handler(_request: httpx.Request) -> httpx.Response: + return build_stream_response( + LLMResultChunk( + model="demo-model", + delta=LLMResultChunkDelta( + index=0, + message=AssistantPromptMessage( + tool_calls=[ + AssistantPromptMessage.ToolCall( + id=call_id, + type="function", + function=AssistantPromptMessage.ToolCall.ToolCallFunction( + name="lookup", arguments="{}" + ), + ) + for call_id in call_ids + ] + ), + ), + ) + ) + + async with self.mock_gateway_stream(httpx.MockTransport(handler)): + adapter = DifyLLMAdapterModel("demo-model", self.make_provider(), model_provider="openai") + async with adapter.request_stream( + [ModelRequest(parts=[UserPromptPart("hello")])], + model_settings=None, + model_request_parameters=ModelRequestParameters(), + ) as stream: + with self.assertRaises(ValueError): + _ = [event async for event in stream] + self.assertEqual(stream.get().parts, []) + async def test_request_splits_embedded_thinking_tags_into_parts(self) -> None: def handler(_request: httpx.Request) -> httpx.Response: return build_stream_response(*single_text_chunk("beforereasoningafter")) diff --git a/dify-agent/uv.lock b/dify-agent/uv.lock index 10578de08f5eb9..b497264c86f764 100644 --- a/dify-agent/uv.lock +++ b/dify-agent/uv.lock @@ -580,7 +580,7 @@ docs = [ requires-dist = [ { name = "e2b", marker = "extra == 'server'", specifier = ">=2.55.1,<3.0.0" }, { name = "fastapi", marker = "extra == 'server'", specifier = "==0.136.0" }, - { name = "graphon", marker = "extra == 'server'", git = "https://github.com/langgenius/graphon?rev=bfeca97726f91eab257e513357b72d456de851b2" }, + { name = "graphon", marker = "extra == 'server'", git = "https://github.com/langgenius/graphon?rev=86f1e11b199bb8ad9111ed49d42edea82a16b7fa" }, { name = "grpcio", marker = "extra == 'server'", specifier = ">=1.60.0,<2.0.0" }, { name = "httpx", specifier = "==0.28.1" }, { name = "httpx2", specifier = ">=2.5.0,<3.0.0" }, @@ -814,7 +814,7 @@ wheels = [ [[package]] name = "graphon" version = "0.8.0" -source = { git = "https://github.com/langgenius/graphon?rev=bfeca97726f91eab257e513357b72d456de851b2#bfeca97726f91eab257e513357b72d456de851b2" } +source = { git = "https://github.com/langgenius/graphon?rev=86f1e11b199bb8ad9111ed49d42edea82a16b7fa#86f1e11b199bb8ad9111ed49d42edea82a16b7fa" } dependencies = [ { name = "charset-normalizer" }, { name = "defusedxml" }, From 02baf0b4e5697cf52166e5643526e57dd7465861 Mon Sep 17 00:00:00 2001 From: WH-2099 Date: Wed, 28 Oct 2026 06:17:17 +0000 Subject: [PATCH 3/4] fix(agent): align replay migration and fixtures with main Keep a single migration head after the current KnowledgeFS chain, emit complete calls once in the model fixture, and update history fixtures and their types for complete database turns. --- .../e9c48a2701bd_add_agent_model_messages.py | 4 +-- .../test_knowledge_fs_migration_chain.py | 12 +++---- .../fixtures/plugin_model.py | 1 - .../core/agent/test_message_replay.py | 36 ++++++++++++++----- .../test_agent_history_prompt_transform.py | 2 +- .../nodes/agent/test_message_history.py | 5 +-- ...t_message_suggested_questions_generator.py | 18 +++++----- ...est_message_suggested_questions_service.py | 16 ++++++--- 8 files changed, 61 insertions(+), 33 deletions(-) diff --git a/api/migrations/versions/e9c48a2701bd_add_agent_model_messages.py b/api/migrations/versions/e9c48a2701bd_add_agent_model_messages.py index 055a7dd99ef77b..f1782234c440a4 100644 --- a/api/migrations/versions/e9c48a2701bd_add_agent_model_messages.py +++ b/api/migrations/versions/e9c48a2701bd_add_agent_model_messages.py @@ -1,7 +1,7 @@ """Persist complete agent model messages. Revision ID: e9c48a2701bd -Revises: a8c3e7f4b912 +Revises: b9d3f7a5c102 """ import sqlalchemy as sa @@ -10,7 +10,7 @@ import models.types revision = "e9c48a2701bd" -down_revision = "a8c3e7f4b912" +down_revision = "b9d3f7a5c102" branch_labels = None depends_on = None diff --git a/api/tests/e2e_tests/contracts/migrations/test_knowledge_fs_migration_chain.py b/api/tests/e2e_tests/contracts/migrations/test_knowledge_fs_migration_chain.py index 4d6a6c3bf22139..752ee1ca021cb6 100644 --- a/api/tests/e2e_tests/contracts/migrations/test_knowledge_fs_migration_chain.py +++ b/api/tests/e2e_tests/contracts/migrations/test_knowledge_fs_migration_chain.py @@ -30,17 +30,17 @@ def _scripts() -> ScriptDirectory: return ScriptDirectory(str(_VERSIONS.parent)) -def test_knowledge_fs_chain_ends_in_the_single_head() -> None: - assert _scripts().get_heads() == [_KNOWLEDGE_FS_CHAIN[-1]] +def test_knowledge_fs_chain_precedes_the_single_head() -> None: + scripts = _scripts() + heads = scripts.get_heads() + assert len(heads) == 1 + assert _KNOWLEDGE_FS_CHAIN[-1] in {revision.revision for revision in scripts.iterate_revisions(heads[0], "base")} def test_upgrade_from_main_head_applies_knowledge_fs_chain_linearly() -> None: scripts = _scripts() - applied_on_main = {revision.revision for revision in scripts.iterate_revisions(_MAIN_HEAD, "base")} pending = [ - revision.revision - for revision in reversed(list(scripts.iterate_revisions("heads", "base"))) - if revision.revision not in applied_on_main + revision.revision for revision in reversed(list(scripts.iterate_revisions(_KNOWLEDGE_FS_CHAIN[-1], _MAIN_HEAD))) ] assert pending == _KNOWLEDGE_FS_CHAIN diff --git a/api/tests/helpers/unit/core/workflow/nodes/parameter_extractor/fixtures/plugin_model.py b/api/tests/helpers/unit/core/workflow/nodes/parameter_extractor/fixtures/plugin_model.py index c3070af4f16416..5e499d23aa841d 100644 --- a/api/tests/helpers/unit/core/workflow/nodes/parameter_extractor/fixtures/plugin_model.py +++ b/api/tests/helpers/unit/core/workflow/nodes/parameter_extractor/fixtures/plugin_model.py @@ -222,7 +222,6 @@ def mocked_chat_create_stream( index=0, message=AssistantPromptMessage( content=full_text[i], - tool_calls=[tool_call] if tool_call else [], ), usage=LLMUsage( prompt_tokens=2, diff --git a/api/tests/unit_tests/core/agent/test_message_replay.py b/api/tests/unit_tests/core/agent/test_message_replay.py index ffae756554ef01..12c64766344219 100644 --- a/api/tests/unit_tests/core/agent/test_message_replay.py +++ b/api/tests/unit_tests/core/agent/test_message_replay.py @@ -1,10 +1,16 @@ +from collections.abc import Generator from copy import deepcopy +from typing import Literal import pytest from sqlalchemy import select +from sqlalchemy.orm import Session from core.agent.cot_chat_agent_runner import CotChatAgentRunner from core.agent.fc_agent_runner import FunctionCallAgentRunner +from core.app.apps.agent_chat.app_config_manager import AgentChatAppConfig +from core.app.apps.message_based_app_queue_manager import MessageBasedAppQueueManager +from core.app.entities.app_invoke_entities import AgentChatAppGenerateEntity from core.model_manager import ModelInstance from graphon.model_runtime.entities import ( AssistantPromptMessage, @@ -12,10 +18,11 @@ LLMResultChunk, LLMResultChunkDelta, LLMUsage, + PromptMessage, TextPromptMessageContent, ToolPromptMessage, ) -from models.model import MessageAgentThought +from models.model import Conversation, Message, MessageAgentThought from services.agent.message_history import AgentMessageHistory from services.workflow_run_agg import WorkflowRunAgg @@ -23,10 +30,17 @@ @pytest.mark.usefixtures("agent_runtime_app") @pytest.mark.parametrize("mode", ["blocking", "streaming", "react"]) def test_tool_rounds_and_persisted_history_preserve_model_messages( - mode, sqlite_session, agent_records, agent_generate_entity, agent_queue_manager, monkeypatch -): + mode: Literal["blocking", "streaming", "react"], + sqlite_session: Session, + agent_records: tuple[Conversation, Message], + agent_generate_entity: AgentChatAppGenerateEntity, + agent_queue_manager: MessageBasedAppQueueManager, + monkeypatch: pytest.MonkeyPatch, +) -> None: conversation, message = agent_records app_config = agent_generate_entity.app_config + assert isinstance(app_config, AgentChatAppConfig) + assert app_config.agent is not None app_config.agent.max_iteration = 3 runner_class = CotChatAgentRunner if mode == "react" else FunctionCallAgentRunner runner = runner_class( @@ -71,10 +85,10 @@ def test_tool_rounds_and_persisted_history_preserve_model_messages( replies.append( AssistantPromptMessage(content='{"action":"Final Answer","action_input":"done"}' if mode == "react" else "done") ) - requests = [] + requests: list[list[PromptMessage]] = [] - def invoke(**kwargs): - requests.append(deepcopy(kwargs["prompt_messages"])) + def invoke(*, prompt_messages: list[PromptMessage], **_kwargs: object) -> LLMResult | Generator[LLMResultChunk]: + requests.append(deepcopy(prompt_messages)) reply = replies[len(requests) - 1].model_copy(deep=True) if mode == "blocking": return LLMResult(model="m", prompt_messages=[], message=reply, usage=LLMUsage.empty_usage()) @@ -88,7 +102,10 @@ def invoke(**kwargs): assert assistants == replies[:2] rows = sqlite_session.scalars(select(MessageAgentThought).order_by(MessageAgentThought.position)).all() - restored = [item for row in rows for item in AgentMessageHistory.model_validate_json(row.model_messages).messages] + restored: list[PromptMessage] = [] + for row in rows: + assert row.model_messages is not None + restored.extend(AgentMessageHistory.model_validate_json(row.model_messages).messages) assert [item for item in restored if isinstance(item, AssistantPromptMessage)] == replies if mode != "react": assert [item.tool_call_id for item in restored if isinstance(item, ToolPromptMessage)] == [ @@ -101,7 +118,7 @@ def invoke(**kwargs): assert [item for item in history if isinstance(item, AssistantPromptMessage)] == replies -def test_history_keeps_empty_tool_results_and_completes_interrupted_calls(): +def test_history_keeps_empty_tool_results_and_completes_interrupted_calls() -> None: assistant = AssistantPromptMessage( opaque_body=False, tool_calls=[ @@ -114,8 +131,11 @@ def test_history_keeps_empty_tool_results_and_completes_interrupted_calls(): ) history = AgentMessageHistory(provider="p", model="m", messages=[assistant]) restored = AgentMessageHistory.model_validate_json(history.model_dump_json()).restore(provider="p", model="m") + assert restored is not None assert restored[0] == assistant + assert isinstance(restored[1], ToolPromptMessage) assert restored[1].tool_call_id == "original" + assert isinstance(restored[1].content, str) assert "interrupted" in restored[1].content history.messages.append(ToolPromptMessage(tool_call_id="original", content="")) assert history.restore(provider="p", model="m") == history.messages diff --git a/api/tests/unit_tests/core/prompt/test_agent_history_prompt_transform.py b/api/tests/unit_tests/core/prompt/test_agent_history_prompt_transform.py index 264b3749ebff00..a37d9ac2d6417c 100644 --- a/api/tests/unit_tests/core/prompt/test_agent_history_prompt_transform.py +++ b/api/tests/unit_tests/core/prompt/test_agent_history_prompt_transform.py @@ -68,7 +68,7 @@ def side_effect_get_num_tokens(*args: object) -> int: @pytest.mark.parametrize("observation_name", ["agent_observation", "knowledge_retrieval"]) def test_history_limit_keeps_complete_agent_turns(mocker: MockerFixture, observation_name: str) -> None: - history = [ + history: list[PromptMessage] = [ UserPromptMessage(content="Earlier query"), AssistantPromptMessage(content="Earlier thought", opaque_body={"signature": "signed"}), UserPromptMessage(content="Tool observation", name=observation_name), diff --git a/api/tests/unit_tests/core/workflow/nodes/agent/test_message_history.py b/api/tests/unit_tests/core/workflow/nodes/agent/test_message_history.py index 1b4f67ff6fa0a7..4a15917dad07d6 100644 --- a/api/tests/unit_tests/core/workflow/nodes/agent/test_message_history.py +++ b/api/tests/unit_tests/core/workflow/nodes/agent/test_message_history.py @@ -1,6 +1,7 @@ import json from collections.abc import Callable, Sequence from datetime import timedelta +from decimal import Decimal from typing import Literal import pytest @@ -192,10 +193,10 @@ def test_workflow_history_budget_preserves_database_turns( _inputs={}, query="latest query", message={}, - message_unit_price=0, + message_unit_price=Decimal(0), answer="latest answer", answer_tokens=1, - answer_unit_price=0, + answer_unit_price=Decimal(0), currency="USD", from_source=older.from_source, ) diff --git a/api/tests/unit_tests/services/test_message_suggested_questions_generator.py b/api/tests/unit_tests/services/test_message_suggested_questions_generator.py index 96846ef7c32bd7..1e68bd3f9437e6 100644 --- a/api/tests/unit_tests/services/test_message_suggested_questions_generator.py +++ b/api/tests/unit_tests/services/test_message_suggested_questions_generator.py @@ -142,13 +142,15 @@ def test_token_transport_failure_propagates_before_best_effort_generation( # request or a replaced transport. The real plugin client wraps that error. monkeypatch.setattr(plugin_base, "plugin_daemon_inner_api_baseurl", URL("unsupported://plugin-daemon")) history = PreparedHistory( - prompts=( - HistoryPrompt( - text="How does this work?", - is_user_message=True, - files=(), - tenant_id=str(uuid4()), - image_detail=ImagePromptMessageContent.DETAIL.HIGH, + turns=( + ( + HistoryPrompt( + text="How does this work?", + is_user_message=True, + files=(), + tenant_id=str(uuid4()), + image_detail=ImagePromptMessageContent.DETAIL.HIGH, + ), ), ) ) @@ -264,7 +266,7 @@ def record_session(session: Session, _transaction: SessionTransaction, _connecti # stored config before it could create a timer or dispatch telemetry. with pytest.raises(ValidationError), current_app.app_context(): SuggestedQuestionsGenerator().generate( - PreparedHistory(prompts=()), + PreparedHistory(turns=()), history_model=make_model_instance(provider="invalid/provider", model="unused"), context=context, instruction_prompt=None, diff --git a/api/tests/unit_tests/services/test_message_suggested_questions_service.py b/api/tests/unit_tests/services/test_message_suggested_questions_service.py index f1f291a7e89607..9dff9b535596c2 100644 --- a/api/tests/unit_tests/services/test_message_suggested_questions_service.py +++ b/api/tests/unit_tests/services/test_message_suggested_questions_service.py @@ -16,7 +16,7 @@ from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom from core.app.file_access import FileAccessScope, bind_file_access_scope -from core.memory.token_buffer_memory import PreparedHistory +from core.memory.token_buffer_memory import HistoryPrompt, PreparedHistory from extensions.ext_database import db from extensions.ext_storage import storage from extensions.storage.opendal_storage import OpenDALStorage @@ -281,12 +281,16 @@ def test_queries_return_detached_configuration_and_history_without_changing_call harness.assert_closed() history = harness.queries.load_history(context) - assert [(prompt.text, prompt.is_user_message) for prompt in history.prompts] == [ + user_prompt, assistant_prompt = history.prompts + assert isinstance(user_prompt, HistoryPrompt) + assert isinstance(assistant_prompt, HistoryPrompt) + prompts = (user_prompt, assistant_prompt) + assert [(prompt.text, prompt.is_user_message) for prompt in prompts] == [ ("How does this work?", True), ("Like this.", False), ] - assert all(not prompt.files for prompt in history.prompts) - assert all(prompt.tenant_id == harness.target.tenant_id for prompt in history.prompts) + assert all(not prompt.files for prompt in prompts) + assert all(prompt.tenant_id == harness.target.tenant_id for prompt in prompts) assert len(harness.read_sessions) == 2 assert harness.read_sessions[0] is not harness.read_sessions[1] harness.assert_closed() @@ -315,7 +319,7 @@ def test_queries_return_detached_configuration_and_history_without_changing_call ) session.execute(update(Message).where(Message.id == harness.message.id).values(answer="Changed answer")) assert context.config == feature - assert history.prompts[-1].text == "Like this." + assert assistant_prompt.text == "Like this." @pytest.mark.parametrize("invoke_from", ["explore", "service-api"]) @@ -485,6 +489,8 @@ def test_history_detaches_remote_attachment_references_before_rendering(harness: harness.assert_closed() assert len(harness.read_sessions) == 2 user_prompt, assistant_prompt = history.prompts + assert isinstance(user_prompt, HistoryPrompt) + assert isinstance(assistant_prompt, HistoryPrompt) assert user_prompt.text == "How does this work?" assert user_prompt.tenant_id == harness.target.tenant_id assert len(user_prompt.files) == 1 From 11eff18013749d4d3a8811d2f20f06b5d5c39437 Mon Sep 17 00:00:00 2001 From: WH-2099 Date: Wed, 28 Oct 2026 06:33:39 +0000 Subject: [PATCH 4/4] test(agent): cover invalid replay history in CI Include the existing replay tests in the KnowledgeFS glue gate and verify that a malformed persisted snapshot falls back atomically to legacy tool history without leaking stale replay data. --- .github/workflows/knowledge-fs-ci.yml | 2 + .../core/agent/test_message_replay.py | 80 +++++++++++++++++++ 2 files changed, 82 insertions(+) diff --git a/.github/workflows/knowledge-fs-ci.yml b/.github/workflows/knowledge-fs-ci.yml index 2e1740d1c6e545..29bc34aaffd7c9 100644 --- a/.github/workflows/knowledge-fs-ci.yml +++ b/.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 diff --git a/api/tests/unit_tests/core/agent/test_message_replay.py b/api/tests/unit_tests/core/agent/test_message_replay.py index 12c64766344219..ba96f76c6d47f3 100644 --- a/api/tests/unit_tests/core/agent/test_message_replay.py +++ b/api/tests/unit_tests/core/agent/test_message_replay.py @@ -1,3 +1,5 @@ +import json +import logging from collections.abc import Generator from copy import deepcopy from typing import Literal @@ -21,7 +23,9 @@ PromptMessage, TextPromptMessageContent, ToolPromptMessage, + UserPromptMessage, ) +from models.enums import CreatorUserRole from models.model import Conversation, Message, MessageAgentThought from services.agent.message_history import AgentMessageHistory from services.workflow_run_agg import WorkflowRunAgg @@ -118,6 +122,82 @@ def invoke(*, prompt_messages: list[PromptMessage], **_kwargs: object) -> LLMRes assert [item for item in history if isinstance(item, AssistantPromptMessage)] == replies +@pytest.mark.usefixtures("agent_runtime_app") +def test_invalid_persisted_model_history_falls_back_to_legacy_tool_turn( + sqlite_session: Session, + agent_records: tuple[Conversation, Message], + agent_generate_entity: AgentChatAppGenerateEntity, + agent_queue_manager: MessageBasedAppQueueManager, + caplog: pytest.LogCaptureFixture, +) -> None: + conversation, message = agent_records + app_config = agent_generate_entity.app_config + assert isinstance(app_config, AgentChatAppConfig) + assert app_config.agent is not None + thought = MessageAgentThought( + message_id=message.id, + position=1, + created_by_role=CreatorUserRole.ACCOUNT, + created_by=agent_generate_entity.user_id, + thought="Checking the saved result.", + tool="lookup", + tool_input=json.dumps({"lookup": {"city": "Boston"}}), + observation=json.dumps({"lookup": "sunny"}), + model_messages=json.dumps( + { + "provider": agent_generate_entity.model_conf.provider, + "model": agent_generate_entity.model_conf.model, + "messages": [ + {"role": "assistant", "content": "Discard this", "opaque_body": {"signature": "stale"}}, + {"role": "invalid", "content": "Corrupt message"}, + ], + } + ), + ) + thought_id = thought.id + sqlite_session.add(thought) + sqlite_session.commit() + next_message = deepcopy(message) + next_message.id = "next-message" + sqlite_session.expunge_all() + + with caplog.at_level(logging.WARNING, logger="core.agent.base_agent_runner"): + runner = FunctionCallAgentRunner( + execution_driver=WorkflowRunAgg.run, + session=sqlite_session, + tenant_id=app_config.tenant_id, + application_generate_entity=agent_generate_entity, + conversation=conversation, + app_config=app_config, + model_config=agent_generate_entity.model_conf, + config=app_config.agent, + queue_manager=agent_queue_manager, + message=next_message, + user_id=agent_generate_entity.user_id, + model_instance=ModelInstance( + provider_model_bundle=agent_generate_entity.model_conf.provider_model_bundle, + model=agent_generate_entity.model_conf.model, + ), + ) + + history = runner.history_prompt_messages + user, assistant, observation = history + assert user == UserPromptMessage(content=message.query) + assert isinstance(assistant, AssistantPromptMessage) + assert assistant.content == "Checking the saved result." + (tool_call,) = assistant.tool_calls + assert tool_call.function.name == "lookup" + assert json.loads(tool_call.function.arguments) == {"city": "Boston"} + assert observation == ToolPromptMessage(tool_call_id=tool_call.id, name="lookup", content="sunny") + assert assistant.opaque_body is None + assert any( + record.name == "core.agent.base_agent_runner" + and record.levelno == logging.WARNING + and thought_id in record.getMessage() + for record in caplog.records + ) + + def test_history_keeps_empty_tool_results_and_completes_interrupted_calls() -> None: assistant = AssistantPromptMessage( opaque_body=False,