From 0d3aab59018f733e6a44476f1d0737718824f511 Mon Sep 17 00:00:00 2001 From: -LAN- Date: Sat, 28 Feb 2026 18:02:39 +0800 Subject: [PATCH] refactor(api): move TokenBufferMemory to model_runtime --- api/.importlinter | 12 ++++++++++++ api/core/agent/base_agent_runner.py | 2 +- api/core/app/apps/agent_chat/app_runner.py | 2 +- api/core/app/apps/base_app_runner.py | 2 +- api/core/app/apps/chat/app_runner.py | 2 +- .../{memory => model_runtime}/token_buffer_memory.py | 0 api/core/prompt/advanced_prompt_transform.py | 2 +- api/core/prompt/agent_history_prompt_transform.py | 2 +- api/core/prompt/prompt_transform.py | 2 +- api/core/prompt/simple_prompt_transform.py | 2 +- api/core/rag/retrieval/dataset_retrieval.py | 2 +- api/core/workflow/nodes/agent/agent_node.py | 2 +- api/core/workflow/nodes/llm/llm_utils.py | 2 +- api/core/workflow/nodes/llm/node.py | 2 +- .../parameter_extractor/parameter_extractor_node.py | 2 +- .../question_classifier/question_classifier_node.py | 2 +- api/services/message_service.py | 2 +- .../core/prompt/test_advanced_prompt_transform.py | 2 +- .../prompt/test_agent_history_prompt_transform.py | 2 +- .../core/prompt/test_simple_prompt_transform.py | 2 +- 20 files changed, 30 insertions(+), 18 deletions(-) rename api/core/{memory => model_runtime}/token_buffer_memory.py (100%) diff --git a/api/.importlinter b/api/.importlinter index 725999c28e..f17546d6a1 100644 --- a/api/.importlinter +++ b/api/.importlinter @@ -247,6 +247,11 @@ ignore_imports = core.workflow.workflow_entry -> models.enums core.workflow.nodes.agent.agent_node -> services core.workflow.nodes.tool.tool_node -> services + core.workflow.nodes.agent.agent_node -> core.model_runtime.token_buffer_memory + core.workflow.nodes.llm.llm_utils -> core.model_runtime.token_buffer_memory + core.workflow.nodes.llm.node -> core.model_runtime.token_buffer_memory + core.workflow.nodes.parameter_extractor.parameter_extractor_node -> core.model_runtime.token_buffer_memory + core.workflow.nodes.question_classifier.question_classifier_node -> core.model_runtime.token_buffer_memory [importlinter:contract:model-runtime-no-internal-imports] name = Model Runtime Internal Imports @@ -299,6 +304,13 @@ ignore_imports = core.model_runtime.model_providers.model_provider_factory -> configs core.model_runtime.model_providers.model_provider_factory -> extensions.ext_redis core.model_runtime.model_providers.model_provider_factory -> models.provider_ids + core.model_runtime.token_buffer_memory -> core.app.app_config.features.file_upload.manager + core.model_runtime.token_buffer_memory -> core.model_manager + core.model_runtime.token_buffer_memory -> core.prompt.utils.extract_thread_messages + core.model_runtime.token_buffer_memory -> core.workflow.file.file_manager + core.model_runtime.token_buffer_memory -> extensions.ext_database + core.model_runtime.token_buffer_memory -> models.model + core.model_runtime.token_buffer_memory -> models.workflow [importlinter:contract:rsc] name = RSC diff --git a/api/core/agent/base_agent_runner.py b/api/core/agent/base_agent_runner.py index 80e180ce96..a1a3ebe8dd 100644 --- a/api/core/agent/base_agent_runner.py +++ b/api/core/agent/base_agent_runner.py @@ -17,7 +17,6 @@ from core.app.entities.app_invoke_entities import ( ) from core.callback_handler.agent_tool_callback_handler import DifyAgentCallbackHandler from core.callback_handler.index_tool_callback_handler import DatasetIndexToolCallbackHandler -from core.memory.token_buffer_memory import TokenBufferMemory from core.model_manager import ModelInstance from core.model_runtime.entities import ( AssistantPromptMessage, @@ -32,6 +31,7 @@ from core.model_runtime.entities import ( from core.model_runtime.entities.message_entities import ImagePromptMessageContent, PromptMessageContentUnionTypes from core.model_runtime.entities.model_entities import ModelFeature from core.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel +from core.model_runtime.token_buffer_memory import TokenBufferMemory from core.prompt.utils.extract_thread_messages import extract_thread_messages from core.tools.__base.tool import Tool from core.tools.entities.tool_entities import ( diff --git a/api/core/app/apps/agent_chat/app_runner.py b/api/core/app/apps/agent_chat/app_runner.py index 7309113f27..df1b475ee2 100644 --- a/api/core/app/apps/agent_chat/app_runner.py +++ b/api/core/app/apps/agent_chat/app_runner.py @@ -12,11 +12,11 @@ from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom from core.app.apps.base_app_runner import AppRunner from core.app.entities.app_invoke_entities import AgentChatAppGenerateEntity from core.app.entities.queue_entities import QueueAnnotationReplyEvent -from core.memory.token_buffer_memory import TokenBufferMemory from core.model_manager import ModelInstance from core.model_runtime.entities.llm_entities import LLMMode from core.model_runtime.entities.model_entities import ModelFeature, ModelPropertyKey from core.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel +from core.model_runtime.token_buffer_memory import TokenBufferMemory from core.moderation.base import ModerationError from extensions.ext_database import db from models.model import App, Conversation, Message diff --git a/api/core/app/apps/base_app_runner.py b/api/core/app/apps/base_app_runner.py index b98e85dbe9..94f81ca0b8 100644 --- a/api/core/app/apps/base_app_runner.py +++ b/api/core/app/apps/base_app_runner.py @@ -22,7 +22,6 @@ from core.app.entities.queue_entities import ( from core.app.features.annotation_reply.annotation_reply import AnnotationReplyFeature from core.app.features.hosting_moderation.hosting_moderation import HostingModerationFeature from core.external_data_tool.external_data_fetch import ExternalDataFetch -from core.memory.token_buffer_memory import TokenBufferMemory from core.model_manager import ModelInstance from core.model_runtime.entities.llm_entities import LLMResult, LLMResultChunk, LLMResultChunkDelta, LLMUsage from core.model_runtime.entities.message_entities import ( @@ -33,6 +32,7 @@ from core.model_runtime.entities.message_entities import ( ) from core.model_runtime.entities.model_entities import ModelPropertyKey from core.model_runtime.errors.invoke import InvokeBadRequestError +from core.model_runtime.token_buffer_memory import TokenBufferMemory from core.moderation.input_moderation import InputModeration from core.prompt.advanced_prompt_transform import AdvancedPromptTransform from core.prompt.entities.advanced_prompt_entities import ChatModelMessage, CompletionModelPromptTemplate, MemoryConfig diff --git a/api/core/app/apps/chat/app_runner.py b/api/core/app/apps/chat/app_runner.py index 4870a56281..bddc778475 100644 --- a/api/core/app/apps/chat/app_runner.py +++ b/api/core/app/apps/chat/app_runner.py @@ -11,9 +11,9 @@ from core.app.entities.app_invoke_entities import ( ) from core.app.entities.queue_entities import QueueAnnotationReplyEvent from core.callback_handler.index_tool_callback_handler import DatasetIndexToolCallbackHandler -from core.memory.token_buffer_memory import TokenBufferMemory from core.model_manager import ModelInstance from core.model_runtime.entities.message_entities import ImagePromptMessageContent +from core.model_runtime.token_buffer_memory import TokenBufferMemory from core.moderation.base import ModerationError from core.rag.retrieval.dataset_retrieval import DatasetRetrieval from core.workflow.file import File diff --git a/api/core/memory/token_buffer_memory.py b/api/core/model_runtime/token_buffer_memory.py similarity index 100% rename from api/core/memory/token_buffer_memory.py rename to api/core/model_runtime/token_buffer_memory.py diff --git a/api/core/prompt/advanced_prompt_transform.py b/api/core/prompt/advanced_prompt_transform.py index fd1b7d838c..7642855b95 100644 --- a/api/core/prompt/advanced_prompt_transform.py +++ b/api/core/prompt/advanced_prompt_transform.py @@ -3,7 +3,6 @@ from typing import cast from core.app.entities.app_invoke_entities import ModelConfigWithCredentialsEntity from core.helper.code_executor.jinja2.jinja2_formatter import Jinja2Formatter -from core.memory.token_buffer_memory import TokenBufferMemory from core.model_runtime.entities import ( AssistantPromptMessage, PromptMessage, @@ -13,6 +12,7 @@ from core.model_runtime.entities import ( UserPromptMessage, ) from core.model_runtime.entities.message_entities import ImagePromptMessageContent, PromptMessageContentUnionTypes +from core.model_runtime.token_buffer_memory import TokenBufferMemory from core.prompt.entities.advanced_prompt_entities import ChatModelMessage, CompletionModelPromptTemplate, MemoryConfig from core.prompt.prompt_transform import PromptTransform from core.prompt.utils.prompt_template_parser import PromptTemplateParser diff --git a/api/core/prompt/agent_history_prompt_transform.py b/api/core/prompt/agent_history_prompt_transform.py index 2b32062140..5039f9553a 100644 --- a/api/core/prompt/agent_history_prompt_transform.py +++ b/api/core/prompt/agent_history_prompt_transform.py @@ -3,13 +3,13 @@ from typing import cast from core.app.entities.app_invoke_entities import ( ModelConfigWithCredentialsEntity, ) -from core.memory.token_buffer_memory import TokenBufferMemory from core.model_runtime.entities.message_entities import ( PromptMessage, SystemPromptMessage, UserPromptMessage, ) from core.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel +from core.model_runtime.token_buffer_memory import TokenBufferMemory from core.prompt.prompt_transform import PromptTransform diff --git a/api/core/prompt/prompt_transform.py b/api/core/prompt/prompt_transform.py index a6e873d587..dc150a90f0 100644 --- a/api/core/prompt/prompt_transform.py +++ b/api/core/prompt/prompt_transform.py @@ -1,10 +1,10 @@ from typing import Any from core.app.entities.app_invoke_entities import ModelConfigWithCredentialsEntity -from core.memory.token_buffer_memory import TokenBufferMemory from core.model_manager import ModelInstance from core.model_runtime.entities.message_entities import PromptMessage from core.model_runtime.entities.model_entities import ModelPropertyKey +from core.model_runtime.token_buffer_memory import TokenBufferMemory from core.prompt.entities.advanced_prompt_entities import MemoryConfig diff --git a/api/core/prompt/simple_prompt_transform.py b/api/core/prompt/simple_prompt_transform.py index d6abbaaa69..7bb5011f12 100644 --- a/api/core/prompt/simple_prompt_transform.py +++ b/api/core/prompt/simple_prompt_transform.py @@ -6,7 +6,6 @@ from typing import TYPE_CHECKING, Any, cast from core.app.app_config.entities import PromptTemplateEntity from core.app.entities.app_invoke_entities import ModelConfigWithCredentialsEntity -from core.memory.token_buffer_memory import TokenBufferMemory from core.model_runtime.entities.message_entities import ( ImagePromptMessageContent, PromptMessage, @@ -15,6 +14,7 @@ from core.model_runtime.entities.message_entities import ( TextPromptMessageContent, UserPromptMessage, ) +from core.model_runtime.token_buffer_memory import TokenBufferMemory from core.prompt.entities.advanced_prompt_entities import MemoryConfig from core.prompt.prompt_transform import PromptTransform from core.prompt.utils.prompt_template_parser import PromptTemplateParser diff --git a/api/core/rag/retrieval/dataset_retrieval.py b/api/core/rag/retrieval/dataset_retrieval.py index cfea8d114a..9e51d4cc5f 100644 --- a/api/core/rag/retrieval/dataset_retrieval.py +++ b/api/core/rag/retrieval/dataset_retrieval.py @@ -23,12 +23,12 @@ from core.callback_handler.index_tool_callback_handler import DatasetIndexToolCa from core.db.session_factory import session_factory from core.entities.agent_entities import PlanningStrategy from core.entities.model_entities import ModelStatus -from core.memory.token_buffer_memory import TokenBufferMemory from core.model_manager import ModelInstance, ModelManager from core.model_runtime.entities.llm_entities import LLMResult, LLMUsage from core.model_runtime.entities.message_entities import PromptMessage, PromptMessageRole, PromptMessageTool from core.model_runtime.entities.model_entities import ModelFeature, ModelType from core.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel +from core.model_runtime.token_buffer_memory import TokenBufferMemory from core.ops.entities.trace_entity import TraceTaskName from core.ops.ops_trace_manager import TraceQueueManager, TraceTask from core.ops.utils import measure_time diff --git a/api/core/workflow/nodes/agent/agent_node.py b/api/core/workflow/nodes/agent/agent_node.py index 5c39a67102..7055dbdb05 100644 --- a/api/core/workflow/nodes/agent/agent_node.py +++ b/api/core/workflow/nodes/agent/agent_node.py @@ -11,10 +11,10 @@ from sqlalchemy.orm import Session from core.agent.entities import AgentToolEntity from core.agent.plugin_entities import AgentStrategyParameter -from core.memory.token_buffer_memory import TokenBufferMemory from core.model_manager import ModelInstance, ModelManager from core.model_runtime.entities.llm_entities import LLMUsage, LLMUsageMetadata from core.model_runtime.entities.model_entities import AIModelEntity, ModelType +from core.model_runtime.token_buffer_memory import TokenBufferMemory from core.model_runtime.utils.encoders import jsonable_encoder from core.provider_manager import ProviderManager from core.tools.entities.tool_entities import ( diff --git a/api/core/workflow/nodes/llm/llm_utils.py b/api/core/workflow/nodes/llm/llm_utils.py index 341a1c1a4c..750a309940 100644 --- a/api/core/workflow/nodes/llm/llm_utils.py +++ b/api/core/workflow/nodes/llm/llm_utils.py @@ -7,11 +7,11 @@ from sqlalchemy.orm import Session from configs import dify_config from core.app.entities.app_invoke_entities import ModelConfigWithCredentialsEntity from core.entities.provider_entities import ProviderQuotaType, QuotaUnit -from core.memory.token_buffer_memory import TokenBufferMemory from core.model_manager import ModelInstance from core.model_runtime.entities.llm_entities import LLMUsage from core.model_runtime.entities.model_entities import ModelType from core.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel +from core.model_runtime.token_buffer_memory import TokenBufferMemory from core.prompt.entities.advanced_prompt_entities import MemoryConfig from core.variables.segments import ArrayAnySegment, ArrayFileSegment, FileSegment, NoneSegment, StringSegment from core.workflow.enums import SystemVariableKey diff --git a/api/core/workflow/nodes/llm/node.py b/api/core/workflow/nodes/llm/node.py index 0259434d90..e6f877cd80 100644 --- a/api/core/workflow/nodes/llm/node.py +++ b/api/core/workflow/nodes/llm/node.py @@ -15,7 +15,6 @@ from core.app.entities.app_invoke_entities import ModelConfigWithCredentialsEnti from core.helper.code_executor import CodeExecutor, CodeLanguage from core.llm_generator.output_parser.errors import OutputParserError from core.llm_generator.output_parser.structured_output import invoke_llm_with_structured_output -from core.memory.token_buffer_memory import TokenBufferMemory from core.model_manager import ModelInstance from core.model_runtime.entities import ( ImagePromptMessageContent, @@ -39,6 +38,7 @@ from core.model_runtime.entities.message_entities import ( UserPromptMessage, ) from core.model_runtime.entities.model_entities import AIModelEntity, ModelFeature, ModelPropertyKey +from core.model_runtime.token_buffer_memory import TokenBufferMemory from core.model_runtime.utils.encoders import jsonable_encoder from core.prompt.entities.advanced_prompt_entities import CompletionModelPromptTemplate, MemoryConfig from core.prompt.utils.prompt_message_util import PromptMessageUtil diff --git a/api/core/workflow/nodes/parameter_extractor/parameter_extractor_node.py b/api/core/workflow/nodes/parameter_extractor/parameter_extractor_node.py index f549d44efa..497281ef1f 100644 --- a/api/core/workflow/nodes/parameter_extractor/parameter_extractor_node.py +++ b/api/core/workflow/nodes/parameter_extractor/parameter_extractor_node.py @@ -6,7 +6,6 @@ from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, cast from core.app.entities.app_invoke_entities import ModelConfigWithCredentialsEntity -from core.memory.token_buffer_memory import TokenBufferMemory from core.model_manager import ModelInstance from core.model_runtime.entities import ImagePromptMessageContent from core.model_runtime.entities.llm_entities import LLMUsage @@ -20,6 +19,7 @@ from core.model_runtime.entities.message_entities import ( ) from core.model_runtime.entities.model_entities import ModelFeature, ModelPropertyKey from core.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel +from core.model_runtime.token_buffer_memory import TokenBufferMemory from core.model_runtime.utils.encoders import jsonable_encoder from core.prompt.advanced_prompt_transform import AdvancedPromptTransform from core.prompt.entities.advanced_prompt_entities import ChatModelMessage, CompletionModelPromptTemplate diff --git a/api/core/workflow/nodes/question_classifier/question_classifier_node.py b/api/core/workflow/nodes/question_classifier/question_classifier_node.py index 3f41c0d0b7..2a4c1f82da 100644 --- a/api/core/workflow/nodes/question_classifier/question_classifier_node.py +++ b/api/core/workflow/nodes/question_classifier/question_classifier_node.py @@ -4,9 +4,9 @@ from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any from core.app.entities.app_invoke_entities import ModelConfigWithCredentialsEntity -from core.memory.token_buffer_memory import TokenBufferMemory from core.model_manager import ModelInstance from core.model_runtime.entities import LLMUsage, ModelPropertyKey, PromptMessageRole +from core.model_runtime.token_buffer_memory import TokenBufferMemory from core.model_runtime.utils.encoders import jsonable_encoder from core.prompt.advanced_prompt_transform import AdvancedPromptTransform from core.prompt.simple_prompt_transform import ModelMode diff --git a/api/services/message_service.py b/api/services/message_service.py index ce699e79d4..f7385a2662 100644 --- a/api/services/message_service.py +++ b/api/services/message_service.py @@ -7,9 +7,9 @@ from sqlalchemy.orm import sessionmaker from core.app.apps.advanced_chat.app_config_manager import AdvancedChatAppConfigManager from core.app.entities.app_invoke_entities import InvokeFrom from core.llm_generator.llm_generator import LLMGenerator -from core.memory.token_buffer_memory import TokenBufferMemory from core.model_manager import ModelManager from core.model_runtime.entities.model_entities import ModelType +from core.model_runtime.token_buffer_memory import TokenBufferMemory from core.ops.entities.trace_entity import TraceTaskName from core.ops.ops_trace_manager import TraceQueueManager, TraceTask from core.ops.utils import measure_time diff --git a/api/tests/unit_tests/core/prompt/test_advanced_prompt_transform.py b/api/tests/unit_tests/core/prompt/test_advanced_prompt_transform.py index f07e55d534..15fba823c5 100644 --- a/api/tests/unit_tests/core/prompt/test_advanced_prompt_transform.py +++ b/api/tests/unit_tests/core/prompt/test_advanced_prompt_transform.py @@ -4,13 +4,13 @@ import pytest from configs import dify_config from core.app.app_config.entities import ModelConfigEntity -from core.memory.token_buffer_memory import TokenBufferMemory from core.model_runtime.entities.message_entities import ( AssistantPromptMessage, ImagePromptMessageContent, PromptMessageRole, UserPromptMessage, ) +from core.model_runtime.token_buffer_memory import TokenBufferMemory from core.prompt.advanced_prompt_transform import AdvancedPromptTransform from core.prompt.entities.advanced_prompt_entities import ChatModelMessage, CompletionModelPromptTemplate, MemoryConfig from core.prompt.utils.prompt_template_parser import PromptTemplateParser 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 d157a41d2c..8efb18e69b 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 @@ -4,7 +4,6 @@ from core.app.entities.app_invoke_entities import ( ModelConfigWithCredentialsEntity, ) from core.entities.provider_configuration import ProviderModelBundle -from core.memory.token_buffer_memory import TokenBufferMemory from core.model_runtime.entities.message_entities import ( AssistantPromptMessage, SystemPromptMessage, @@ -12,6 +11,7 @@ from core.model_runtime.entities.message_entities import ( UserPromptMessage, ) from core.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel +from core.model_runtime.token_buffer_memory import TokenBufferMemory from core.prompt.agent_history_prompt_transform import AgentHistoryPromptTransform from models.model import Conversation diff --git a/api/tests/unit_tests/core/prompt/test_simple_prompt_transform.py b/api/tests/unit_tests/core/prompt/test_simple_prompt_transform.py index c822ecbe78..e7ebc3cd0d 100644 --- a/api/tests/unit_tests/core/prompt/test_simple_prompt_transform.py +++ b/api/tests/unit_tests/core/prompt/test_simple_prompt_transform.py @@ -1,8 +1,8 @@ from unittest.mock import MagicMock from core.app.entities.app_invoke_entities import ModelConfigWithCredentialsEntity -from core.memory.token_buffer_memory import TokenBufferMemory from core.model_runtime.entities.message_entities import AssistantPromptMessage, UserPromptMessage +from core.model_runtime.token_buffer_memory import TokenBufferMemory from core.prompt.simple_prompt_transform import SimplePromptTransform from models.model import AppMode, Conversation