youngkingdom/dify
1
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-05-04 09:28:04 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor: inject memory interface into LLMNode (#32754)

Browse Source 
Signed-off-by: -LAN- <laipz8200@outlook.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
-LAN-
2026-03-01 04:05:18 +08:00
committed by GitHub
parent 1f0fca89a8
commit c034eb036c
4 changed files with 115 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

39
api/tests/unit_tests/core/workflow/nodes/llm/test_node.py
View File

@ -12,6 +12,7 @@ from core.entities.provider_entities import CustomConfiguration, SystemConfigura
from core.model_manager import ModelInstance
from core.model_runtime.entities.common_entities import I18nObject
from core.model_runtime.entities.message_entities import (
AssistantPromptMessage,
ImagePromptMessageContent,
PromptMessage,
PromptMessageRole,
@ -20,6 +21,7 @@ from core.model_runtime.entities.message_entities import (
)
from core.model_runtime.entities.model_entities import AIModelEntity, FetchFrom, ModelType
from core.model_runtime.model_providers.model_provider_factory import ModelProviderFactory
from core.prompt.entities.advanced_prompt_entities import MemoryConfig
from core.workflow.entities import GraphInitParams
from core.workflow.file import File, FileTransferMethod, FileType
from core.workflow.nodes.llm import llm_utils
@ -32,7 +34,7 @@ from core.workflow.nodes.llm.entities import (
VisionConfigOptions,
)
from core.workflow.nodes.llm.file_saver import LLMFileSaver
from core.workflow.nodes.llm.node import LLMNode
from core.workflow.nodes.llm.node import LLMNode, _handle_memory_completion_mode
from core.workflow.nodes.llm.protocols import CredentialsProvider, ModelFactory
from core.workflow.runtime import GraphRuntimeState, VariablePool
from core.workflow.system_variable import SystemVariable
@ -587,6 +589,41 @@ def test_handle_list_messages_basic(llm_node):
assert result[0].content == [TextPromptMessageContent(data="Hello, world")]
def test_handle_memory_completion_mode_uses_prompt_message_interface():
memory = mock.MagicMock(spec=MockTokenBufferMemory)
memory.get_history_prompt_messages.return_value = [
UserPromptMessage(
content=[
TextPromptMessageContent(data="first question"),
ImagePromptMessageContent(
format="png",
url="https://example.com/image.png",
mime_type="image/png",
),
]
),
AssistantPromptMessage(content="first answer"),
]
model_instance = mock.MagicMock(spec=ModelInstance)
memory_config = MemoryConfig(
role_prefix=MemoryConfig.RolePrefix(user="Human", assistant="Assistant"),
window=MemoryConfig.WindowConfig(enabled=True, size=3),
)
with mock.patch("core.workflow.nodes.llm.node._calculate_rest_token", return_value=2000) as mock_rest_token:
memory_text = _handle_memory_completion_mode(
memory=memory,
memory_config=memory_config,
model_instance=model_instance,
)
assert memory_text == "Human: first question\n[image]\nAssistant: first answer"
mock_rest_token.assert_called_once_with(prompt_messages=[], model_instance=model_instance)
memory.get_history_prompt_messages.assert_called_once_with(max_token_limit=2000, message_limit=3)
@pytest.fixture
def llm_node_for_multimodal(llm_node_data, graph_init_params, graph_runtime_state) -> tuple[LLMNode, LLMFileSaver]:
mock_file_saver: LLMFileSaver = mock.MagicMock(spec=LLMFileSaver)