refactor(api): move TokenBufferMemory to model_runtime

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

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 (

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

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

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

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

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
 
 

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
 
 

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

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

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 (

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

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

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

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

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

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

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
 

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