adai/MemoryBear
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor(core): migrate memory write tasks to centralized scheduler

Browse Source
This commit is contained in:
Eternity
2026-04-22 16:50:06 +08:00
parent 6f323f2435
commit c5ae82c3c2
15 changed files with 358 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files

56
api/app/core/memory/agent/langgraph_graph/routing/write_router.py
View File

@@ -12,9 +12,8 @@ from app.core.memory.utils.llm.llm_utils import MemoryClientFactory
from app.db import get_db_context
from app.repositories.memory_short_repository import LongTermMemoryRepository
from app.schemas.memory_agent_schema import AgentMemory_Long_Term
from app.services.task_service import get_task_memory_write_result
from app.tasks import write_message_task
from app.utils.config_utils import resolve_config_id
from celery_task_scheduler import scheduler
logger = get_agent_logger(__name__)
template_root = os.path.join(PROJECT_ROOT_, 'memory', 'agent', 'utils', 'prompt')
@@ -86,16 +85,28 @@ async def write(
logger.info(
f"[WRITE] Submitting Celery task - user={actual_end_user_id}, messages={len(structured_messages)}, config={actual_config_id}")
write_id = write_message_task.delay(
actual_end_user_id, # end_user_id: User ID
structured_messages, # message: JSON string format message list
str(actual_config_id), # config_id: Configuration ID string
storage_type, # storage_type: "neo4j"
user_rag_memory_id or "" # user_rag_memory_id: RAG memory ID (not used in Neo4j mode)
# write_id = write_message_task.delay(
# actual_end_user_id, # end_user_id: User ID
# structured_messages, # message: JSON string format message list
# str(actual_config_id), # config_id: Configuration ID string
# storage_type, # storage_type: "neo4j"
# user_rag_memory_id or "" # user_rag_memory_id: RAG memory ID (not used in Neo4j mode)
# )
scheduler.push_task(
"app.core.memory.agent.write_message",
end_user_id,
{
"end_user_id": end_user_id,
"message": structured_messages,
"config_id": str(actual_config_id),
"storage_type": storage_type,
"user_rag_memory_id": user_rag_memory_id or ""
}
)
logger.info(f"[WRITE] Celery task submitted - task_id={write_id}")
write_status = get_task_memory_write_result(str(write_id))
logger.info(f'[WRITE] Task result - user={actual_end_user_id}, status={write_status}')
# logger.info(f"[WRITE] Celery task submitted - task_id={write_id}")
# write_status = get_task_memory_write_result(str(write_id))
# logger.info(f'[WRITE] Task result - user={actual_end_user_id}')
async def term_memory_save(end_user_id, strategy_type, scope):
@@ -164,13 +175,24 @@ async def window_dialogue(end_user_id, langchain_messages, memory_config, scope)
else:
config_id = memory_config
write_message_task.delay(
end_user_id, # end_user_id: User ID
redis_messages, # message: JSON string format message list
config_id, # config_id: Configuration ID string
AgentMemory_Long_Term.STORAGE_NEO4J, # storage_type: "neo4j"
"" # user_rag_memory_id: RAG memory ID (not used in Neo4j mode)
scheduler.push_task(
"app.core.memory.agent.write_message",
end_user_id,
{
"end_user_id": end_user_id,
"message": redis_messages,
"config_id": config_id,
"storage_type": AgentMemory_Long_Term.STORAGE_NEO4J,
"user_rag_memory_id": ""
}
)
# write_message_task.delay(
# end_user_id, # end_user_id: User ID
# redis_messages, # message: JSON string format message list
# config_id, # config_id: Configuration ID string
# AgentMemory_Long_Term.STORAGE_NEO4J, # storage_type: "neo4j"
# "" # user_rag_memory_id: RAG memory ID (not used in Neo4j mode)
# )
count_store.update_sessions_count(end_user_id, 0, [])

4
api/app/core/memory/pipelines/memory_read.py
View File

@@ -1,8 +1,8 @@
from app.core.memory.enums import SearchStrategy, StorageType
from app.core.memory.models.service_models import MemorySearchResult
from app.core.memory.pipelines.base_pipeline import ModelClientMixin, DBRequiredPipeline
from app.core.memory.read_services.content_search import Neo4jSearchService, RAGSearchService
from app.core.memory.read_services.query_preprocessor import QueryPreprocessor
from core.memory.read_services.search_engine.content_search import Neo4jSearchService, RAGSearchService
from core.memory.read_services.generate_engine.query_preprocessor import QueryPreprocessor
class ReadPipeLine(ModelClientMixin, DBRequiredPipeline):

0
api/app/core/memory/read_services/generate_engine/__init__.py Normal file
View File

0
api/app/core/memory/read_services/query_preprocessor.py → api/app/core/memory/read_services/generate_engine/query_preprocessor.py
View File

0
api/app/core/memory/read_services/retrieval_summary.py → api/app/core/memory/read_services/generate_engine/retrieval_summary.py
View File

0
api/app/core/memory/read_services/search_engine/__init__.py Normal file
View File

2
api/app/core/memory/read_services/content_search.py → api/app/core/memory/read_services/search_engine/content_search.py
View File

@@ -8,7 +8,7 @@ from neo4j import Session
from app.core.memory.enums import Neo4jNodeType
from app.core.memory.memory_service import MemoryContext
from app.core.memory.models.service_models import Memory, MemorySearchResult
from app.core.memory.read_services.result_builder import data_builder_factory
from core.memory.read_services.search_engine.result_builder import data_builder_factory
from app.core.models import RedBearEmbeddings
from app.core.rag.nlp.search import knowledge_retrieval
from app.repositories import knowledge_repository

0
api/app/core/memory/read_services/result_builder.py → api/app/core/memory/read_services/search_engine/result_builder.py
View File

25
api/app/core/workflow/nodes/memory/node.py
View File

@@ -11,7 +11,7 @@ from app.core.workflow.variable.base_variable import VariableType
from app.core.workflow.variable.variable_objects import FileVariable, ArrayVariable
from app.db import get_db_read
from app.schemas import FileInput
from app.tasks import write_message_task
from celery_task_scheduler import scheduler
class MemoryReadNode(BaseNode):
@@ -126,12 +126,23 @@ class MemoryWriteNode(BaseNode):
"files": file_info
})
write_message_task.delay(
end_user_id=end_user_id,
message=messages,
config_id=str(self.typed_config.config_id),
storage_type=state["memory_storage_type"],
user_rag_memory_id=state["user_rag_memory_id"]
scheduler.push_task(
"app.core.memory.agent.write_message",
end_user_id,
{
"end_user_id": end_user_id,
"message": messages,
"config_id": str(self.typed_config.config_id),
"storage_type": state["memory_storage_type"],
"user_rag_memory_id": state["user_rag_memory_id"]
}
)
# write_message_task.delay(
# end_user_id=end_user_id,
# message=messages,
# config_id=str(self.typed_config.config_id),
# storage_type=state["memory_storage_type"],
# user_rag_memory_id=state["user_rag_memory_id"]
# )
return "success"