import asyncio import json import os import re import time import uuid from uuid import UUID from datetime import datetime, timezone from math import ceil from typing import Any, Dict, List, Optional import redis import requests import trio # Import a unified Celery instance from app.celery_app import celery_app from app.core.config import settings from app.core.rag.graphrag.general.index import init_graphrag, run_graphrag_for_kb from app.core.rag.graphrag.utils import get_llm_cache, set_llm_cache from app.core.rag.llm.chat_model import Base from app.core.rag.llm.cv_model import QWenCV from app.core.rag.llm.embedding_model import OpenAIEmbed from app.core.rag.llm.sequence2txt_model import QWenSeq2txt from app.core.rag.models.chunk import DocumentChunk from app.core.rag.prompts.generator import question_proposal from app.core.rag.vdb.elasticsearch.elasticsearch_vector import ( ElasticSearchVectorFactory, ) from app.db import get_db, get_db_context from app.models.document_model import Document from app.models.knowledge_model import Knowledge from app.services.memory_agent_service import MemoryAgentService @celery_app.task(name="tasks.process_item") def process_item(item: dict): """ A simulated long-running task that processes an item. In a real-world scenario, this could be anything: - Sending an email - Generating a report - Performing a complex calculation - Calling a third-party API """ print(f"Processing item: {item['name']}") # Simulate work for 5 seconds time.sleep(5) result = f"Item '{item['name']}' processed successfully at a price of ${item['price']}." print(result) return result @celery_app.task(name="app.core.rag.tasks.parse_document") def parse_document(file_path: str, document_id: uuid.UUID): """ Document parsing, vectorization, and storage """ db = next(get_db()) # Manually call the generator db_document = None db_knowledge = None progress_msg = f"{datetime.now().strftime('%H:%M:%S')} Task has been received.\n" try: db_document = db.query(Document).filter(Document.id == document_id).first() db_knowledge = db.query(Knowledge).filter(Knowledge.id == db_document.kb_id).first() # 1. Document parsing & segmentation progress_msg += f"{datetime.now().strftime('%H:%M:%S')} Start to parse.\n" start_time = time.time() db_document.progress = 0.0 db_document.progress_msg = progress_msg db_document.process_begin_at = datetime.now(tz=timezone.utc) db_document.process_duration = 0.0 db_document.run = 1 db.commit() db.refresh(db_document) def progress_callback(prog=None, msg=None): nonlocal progress_msg # Declare the use of an external progress_msg variable progress_msg += f"{datetime.now().strftime('%H:%M:%S')} parse progress: {prog} msg: {msg}.\n" # Prepare to configure chat_mdl、embedding_model、vision_model information chat_model = Base( key=db_knowledge.llm.api_keys[0].api_key, model_name=db_knowledge.llm.api_keys[0].model_name, base_url=db_knowledge.llm.api_keys[0].api_base ) embedding_model = OpenAIEmbed( key=db_knowledge.embedding.api_keys[0].api_key, model_name=db_knowledge.embedding.api_keys[0].model_name, base_url=db_knowledge.embedding.api_keys[0].api_base ) vision_model = QWenCV( key=db_knowledge.image2text.api_keys[0].api_key, model_name=db_knowledge.image2text.api_keys[0].model_name, lang="Chinese", base_url=db_knowledge.image2text.api_keys[0].api_base ) if re.search(r"\.(da|wave|wav|mp3|aac|flac|ogg|aiff|au|midi|wma|realaudio|vqf|oggvorbis|ape?)$", file_path, re.IGNORECASE): vision_model = QWenSeq2txt( key=os.getenv("QWEN3_OMNI_API_KEY", ""), model_name=os.getenv("QWEN3_OMNI_MODEL_NAME", "qwen3-omni-flash"), lang="Chinese", base_url=os.getenv("QWEN3_OMNI_BASE_URL", "https://dashscope.aliyuncs.com/compatible-mode/v1"), ) elif re.search(r"\.(png|jpeg|jpg|gif|bmp|svg|mp4|mov|avi|flv|mpeg|mpg|webm|wmv|3gp|3gpp|mkv?)$", file_path, re.IGNORECASE): vision_model = QWenCV( key=os.getenv("QWEN3_OMNI_API_KEY", ""), model_name=os.getenv("QWEN3_OMNI_MODEL_NAME", "qwen3-omni-flash"), lang="Chinese", base_url=os.getenv("QWEN3_OMNI_BASE_URL", "https://dashscope.aliyuncs.com/compatible-mode/v1"), ) else: print(file_path) from app.core.rag.app.naive import chunk res = chunk(filename=file_path, from_page=0, to_page=100000, callback=progress_callback, vision_model=vision_model, parser_config=db_document.parser_config, is_root=False) progress_msg += f"{datetime.now().strftime('%H:%M:%S')} Finish parsing.\n" db_document.progress = 0.8 db_document.progress_msg = progress_msg db.commit() db.refresh(db_document) # 2. Document vectorization and storage total_chunks = len(res) progress_msg += f"{datetime.now().strftime('%H:%M:%S')} Generate {total_chunks} chunks.\n" batch_size = 100 total_batches = ceil(total_chunks / batch_size) progress_per_batch = 0.2 / total_batches # Progress of each batch vector_service = ElasticSearchVectorFactory().init_vector(knowledge=db_knowledge) # 2.1 Delete document vector index vector_service.delete_by_metadata_field(key="document_id", value=str(document_id)) # 2.2 Vectorize and import batch documents for batch_start in range(0, total_chunks, batch_size): batch_end = min(batch_start + batch_size, total_chunks) # prevent out-of-bounds batch = res[batch_start: batch_end] # Retrieve the current batch chunks = [] # Process the current batch for idx_in_batch, item in enumerate(batch): global_idx = batch_start + idx_in_batch # Calculate global index metadata = { "doc_id": uuid.uuid4().hex, "file_id": str(db_document.file_id), "file_name": db_document.file_name, "file_created_at": int(db_document.created_at.timestamp() * 1000), "document_id": str(db_document.id), "knowledge_id": str(db_document.kb_id), "sort_id": global_idx, "status": 1, } if db_document.parser_config.get("auto_questions", 0): topn = db_document.parser_config["auto_questions"] cached = get_llm_cache(chat_model.model_name, item["content_with_weight"], "question", {"topn": topn}) if not cached: cached = question_proposal(chat_model, item["content_with_weight"], topn) set_llm_cache(chat_model.model_name, item["content_with_weight"], cached, "question", {"topn": topn}) chunks.append( DocumentChunk(page_content=f"question: {cached} answer: {item['content_with_weight']}", metadata=metadata)) else: chunks.append(DocumentChunk(page_content=item["content_with_weight"], metadata=metadata)) # Bulk segmented vector import vector_service.add_chunks(chunks) # Update progress db_document.progress += progress_per_batch progress_msg += f"{datetime.now().strftime('%H:%M:%S')} Embedding progress ({db_document.progress}).\n" db_document.progress_msg = progress_msg db_document.process_duration = time.time() - start_time db_document.run = 0 db.commit() db.refresh(db_document) # Vectorization and data entry completed progress_msg += f"{datetime.now().strftime('%H:%M:%S')} Indexing done.\n" db_document.chunk_num = total_chunks db_document.progress = 1.0 db_document.process_duration = time.time() - start_time progress_msg += f"{datetime.now().strftime('%H:%M:%S')} Task done ({db_document.process_duration}s).\n" db_document.progress_msg = progress_msg db_document.run = 0 db.commit() # using graphrag if db_knowledge.parser_config and db_knowledge.parser_config.get("graphrag", {}).get("use_graphrag", False): graphrag_conf = db_knowledge.parser_config.get("graphrag", {}) with_resolution = graphrag_conf.get("resolution", False) with_community = graphrag_conf.get("community", False) def callback(*args, msg=None, **kwargs): nonlocal progress_msg message = msg or (args[0] if args else "No message") progress_msg += f"{datetime.now().strftime('%H:%M:%S')} run graphrag msg: {message}.\n" progress_msg += f"{datetime.now().strftime('%H:%M:%S')} Start to run graphrag.\n" start_time = time.time() db_document.progress_msg = progress_msg db.commit() db.refresh(db_document) task = { "id": str(db_document.id), "workspace_id": str(db_knowledge.workspace_id), "kb_id": str(db_knowledge.id), "parser_config": db_knowledge.parser_config, } # init_graphrag vts, _ = embedding_model.encode(["ok"]) vector_size = len(vts[0]) init_graphrag(task, vector_size) async def _run(row: dict, document_ids: list[str], language: str, parser_config: dict, vector_service, chat_model, embedding_model, callback, with_resolution: bool = True, with_community: bool = True, ) -> dict: await trio.sleep(5) # Delay for 10 seconds nonlocal progress_msg # Declare the use of an external progress_msg variable result = await run_graphrag_for_kb( row=row, document_ids=document_ids, language=language, parser_config=parser_config, vector_service=vector_service, chat_model=chat_model, embedding_model=embedding_model, callback=callback, with_resolution=with_resolution, with_community=with_community, ) progress_msg += f"{datetime.now().strftime('%H:%M:%S')} GraphRAG task result for task {task}:\n{result}\n" return result try: trio.run( lambda: _run( row=task, document_ids=[str(db_document.id)], language="Chinese", parser_config=db_knowledge.parser_config, vector_service=vector_service, chat_model=chat_model, embedding_model=embedding_model, callback=callback, with_resolution=with_resolution, with_community=with_community, ) ) except Exception as e: progress_msg += f"{datetime.now().strftime('%H:%M:%S')} GraphRAG task failed for task {task}:\n{str(e)}\n" progress_msg += f"{datetime.now().strftime('%H:%M:%S')} Knowledge Graph done ({time.time() - start_time}s)" db_document.progress_msg = progress_msg db.commit() db.refresh(db_document) result = f"parse document '{db_document.file_name}' processed successfully." return result except Exception as e: if 'db_document' in locals(): db_document.progress_msg += f"Failed to vectorize and import the parsed document:{str(e)}\n" db_document.run = 0 db.commit() result = f"parse document '{db_document.file_name}' failed." return result finally: db.close() @celery_app.task(name="app.core.rag.tasks.build_graphrag_for_kb") def build_graphrag_for_kb(kb_id: uuid.UUID): """ build knowledge graph """ db = next(get_db()) # Manually call the generator db_documents = None db_knowledge = None try: db_documents = db.query(Document).filter(Document.kb_id == kb_id).all() db_knowledge = db.query(Knowledge).filter(Knowledge.id == kb_id).first() # 1. Prepare to configure chat_mdl、embedding_model、vision_model information chat_model = Base( key=db_knowledge.llm.api_keys[0].api_key, model_name=db_knowledge.llm.api_keys[0].model_name, base_url=db_knowledge.llm.api_keys[0].api_base ) embedding_model = OpenAIEmbed( key=db_knowledge.embedding.api_keys[0].api_key, model_name=db_knowledge.embedding.api_keys[0].model_name, base_url=db_knowledge.embedding.api_keys[0].api_base ) vision_model = QWenCV( key=db_knowledge.image2text.api_keys[0].api_key, model_name=db_knowledge.image2text.api_keys[0].model_name, lang="Chinese", base_url=db_knowledge.image2text.api_keys[0].api_base ) # 2. get all document_ids from knowledge base vector_service = ElasticSearchVectorFactory().init_vector(knowledge=db_knowledge) total, items = vector_service.search_by_segment(document_id=None, query=None, pagesize=9999, page=1, asc=True) document_ids = [str(item.id) for item in db_documents] # 2. using graphrag if db_knowledge.parser_config and db_knowledge.parser_config.get("graphrag", {}).get("use_graphrag", False): graphrag_conf = db_knowledge.parser_config.get("graphrag", {}) with_resolution = graphrag_conf.get("resolution", False) with_community = graphrag_conf.get("community", False) def callback(*args, msg=None, **kwargs): message = msg or (args[0] if args else "No message") print(f"{datetime.now().strftime('%H:%M:%S')} run graphrag msg: {message}.\n") start_time = time.time() task = { "id": str(db_knowledge.id), "workspace_id": str(db_knowledge.workspace_id), "kb_id": str(db_knowledge.id), "parser_config": db_knowledge.parser_config, } # init_graphrag vts, _ = embedding_model.encode(["ok"]) vector_size = len(vts[0]) init_graphrag(task, vector_size) async def _run(row: dict, document_ids: list[str], language: str, parser_config: dict, vector_service, chat_model, embedding_model, callback, with_resolution: bool = True, with_community: bool = True, ) -> dict: result = await run_graphrag_for_kb( row=row, document_ids=document_ids, language=language, parser_config=parser_config, vector_service=vector_service, chat_model=chat_model, embedding_model=embedding_model, callback=callback, with_resolution=with_resolution, with_community=with_community, ) print(f"{datetime.now().strftime('%H:%M:%S')} GraphRAG task result for task {task}:\n{result}\n") return result try: trio.run( lambda: _run( row=task, document_ids=document_ids, language="Chinese", parser_config=db_knowledge.parser_config, vector_service=vector_service, chat_model=chat_model, embedding_model=embedding_model, callback=callback, with_resolution=with_resolution, with_community=with_community, ) ) except Exception as e: print(f"{datetime.now().strftime('%H:%M:%S')} GraphRAG task failed for task {task}:\n{str(e)}\n") print(f"{datetime.now().strftime('%H:%M:%S')} Knowledge Graph done ({time.time() - start_time}s)") result = f"build knowledge graph '{db_knowledge.name}' processed successfully." return result except Exception as e: if 'db_knowledge' in locals(): print(f"Failed to build knowledge grap:{str(e)}\n") result = f"build knowledge grap '{db_knowledge.name}' failed." return result finally: db.close() @celery_app.task(name="app.core.memory.agent.read_message", bind=True) def read_message_task(self, end_user_id: str, message: str, history: List[Dict[str, Any]], search_switch: str, config_id: str, storage_type:str, user_rag_memory_id:str) -> Dict[str, Any]: """Celery task to process a read message via MemoryAgentService. Args: end_user_id: Group ID for the memory agent (also used as end_user_id) message: User message to process history: Conversation history search_switch: Search switch parameter config_id: Configuration ID as string (will be converted to UUID) Returns: Dict containing the result and metadata Raises: Exception on failure """ start_time = time.time() # Convert config_id string to UUID actual_config_id = None if config_id: try: actual_config_id = uuid.UUID(config_id) if isinstance(config_id, str) else config_id except (ValueError, AttributeError): # If conversion fails, leave as None and try to resolve pass # Resolve config_id if None if actual_config_id is None: try: from app.services.memory_agent_service import get_end_user_connected_config db = next(get_db()) try: connected_config = get_end_user_connected_config(end_user_id, db) actual_config_id = connected_config.get("memory_config_id") finally: db.close() except Exception: # Log but continue - will fail later with proper error pass async def _run() -> str: db = next(get_db()) try: service = MemoryAgentService() return await service.read_memory(end_user_id, message, history, search_switch, actual_config_id, db, storage_type, user_rag_memory_id) finally: db.close() try: # 使用 nest_asyncio 来避免事件循环冲突 try: import nest_asyncio nest_asyncio.apply() except ImportError: pass # 尝试获取现有事件循环，如果不存在则创建新的 try: loop = asyncio.get_event_loop() if loop.is_closed(): loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) except RuntimeError: loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) result = loop.run_until_complete(_run()) elapsed_time = time.time() - start_time return { "status": "SUCCESS", "result": result, "end_user_id": end_user_id, "config_id": config_id, "elapsed_time": elapsed_time, "task_id": self.request.id } except BaseException as e: elapsed_time = time.time() - start_time # Handle ExceptionGroup from TaskGroup if hasattr(e, 'exceptions'): error_messages = [f"{type(sub_e).__name__}: {str(sub_e)}" for sub_e in e.exceptions] detailed_error = "; ".join(error_messages) else: detailed_error = str(e) return { "status": "FAILURE", "error": detailed_error, "end_user_id": end_user_id, "config_id": config_id, "elapsed_time": elapsed_time, "task_id": self.request.id } @celery_app.task(name="app.core.memory.agent.write_message", bind=True) def write_message_task(self, end_user_id: str, message: str, config_id: str, storage_type:str, user_rag_memory_id:str) -> Dict[str, Any]: """Celery task to process a write message via MemoryAgentService. Args: end_user_id: Group ID for the memory agent (also used as end_user_id) message: Message to write config_id: Configuration ID as string (will be converted to UUID) Returns: Dict containing the result and metadata Raises: Exception on failure """ from app.core.logging_config import get_logger logger = get_logger(__name__) logger.info(f"[CELERY WRITE] Starting write task - end_user_id={end_user_id}, config_id={config_id}, storage_type={storage_type}") start_time = time.time() # Convert config_id string to UUID actual_config_id = None if config_id: try: actual_config_id = uuid.UUID(config_id) if isinstance(config_id, str) else config_id logger.info(f"[CELERY WRITE] Converted config_id to UUID: {actual_config_id} (type: {type(actual_config_id).__name__})") except (ValueError, AttributeError) as e: logger.error(f"[CELERY WRITE] Invalid config_id format: {config_id}, error: {e}") return { "status": "FAILURE", "error": f"Invalid config_id format: {config_id}", "end_user_id": end_user_id, "config_id": config_id, "elapsed_time": 0.0, "task_id": self.request.id } # Resolve config_id if None if actual_config_id is None: try: from app.services.memory_agent_service import get_end_user_connected_config db = next(get_db()) try: connected_config = get_end_user_connected_config(end_user_id, db) actual_config_id = connected_config.get("memory_config_id") finally: db.close() except Exception: # Log but continue - will fail later with proper error pass async def _run() -> str: db = next(get_db()) try: logger.info(f"[CELERY WRITE] Executing MemoryAgentService.write_memory with config_id={actual_config_id} (type: {type(actual_config_id).__name__})") service = MemoryAgentService() result = await service.write_memory(end_user_id, message, actual_config_id, db, storage_type, user_rag_memory_id) logger.info(f"[CELERY WRITE] Write completed successfully: {result}") return result except Exception as e: logger.error(f"[CELERY WRITE] Write failed: {e}", exc_info=True) raise finally: db.close() try: # 使用 nest_asyncio 来避免事件循环冲突 try: import nest_asyncio nest_asyncio.apply() except ImportError: pass # 尝试获取现有事件循环，如果不存在则创建新的 try: loop = asyncio.get_event_loop() if loop.is_closed(): loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) except RuntimeError: loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) result = loop.run_until_complete(_run()) elapsed_time = time.time() - start_time logger.info(f"[CELERY WRITE] Task completed successfully - elapsed_time={elapsed_time:.2f}s, task_id={self.request.id}") return { "status": "SUCCESS", "result": result, "end_user_id": end_user_id, "config_id": config_id, "elapsed_time": elapsed_time, "task_id": self.request.id } except BaseException as e: elapsed_time = time.time() - start_time # Handle ExceptionGroup from TaskGroup if hasattr(e, 'exceptions'): error_messages = [f"{type(sub_e).__name__}: {str(sub_e)}" for sub_e in e.exceptions] detailed_error = "; ".join(error_messages) else: detailed_error = str(e) logger.error(f"[CELERY WRITE] Task failed - elapsed_time={elapsed_time:.2f}s, error={detailed_error}", exc_info=True) return { "status": "FAILURE", "error": detailed_error, "end_user_id": end_user_id, "config_id": config_id, "elapsed_time": elapsed_time, "task_id": self.request.id } def reflection_engine() -> None: """Empty function placeholder for timed background reflection. Intentionally left blank; replace with real reflection logic later. """ import asyncio from app.core.memory.utils.self_reflexion_utils.self_reflexion import self_reflexion host_id = uuid.UUID("2f6ff1eb-50c7-4765-8e89-e4566be19122") asyncio.run(self_reflexion(host_id)) @celery_app.task(name="app.core.memory.agent.reflection.timer") def reflection_timer_task() -> None: """Periodic Celery task that invokes reflection_engine. Raises an exception on failure. """ reflection_engine() # unused task # @celery_app.task(name="app.core.memory.agent.health.check_read_service") # def check_read_service_task() -> Dict[str, str]: # """Call read_service and write latest status to Redis. # Returns status data dict that gets written to Redis. # """ # client = redis.Redis( # host=settings.REDIS_HOST, # port=settings.REDIS_PORT, # db=settings.REDIS_DB, # password=settings.REDIS_PASSWORD if settings.REDIS_PASSWORD else None # ) # try: # api_url = f"http://{settings.SERVER_IP}:8000/api/memory/read_service" # payload = { # "user_id": "健康检查", # "apply_id": "健康检查", # "group_id": "健康检查", # "message": "你好", # "history": [], # "search_switch": "2", # } # resp = requests.post(api_url, json=payload, timeout=15) # ok = resp.status_code == 200 # status = "Success" if ok else "Fail" # msg = "接口请求成功" if ok else f"接口请求失败: {resp.status_code}" # error = "" if ok else resp.text # code = 0 if ok else 500 # except Exception as e: # status = "Fail" # msg = "接口请求失败" # error = str(e) # code = 500 # data = { # "status": status, # "msg": msg, # "error": error, # "code": str(code), # "time": str(int(time.time())), # } # client.hset("memsci:health:read_service", mapping=data) # client.expire("memsci:health:read_service", int(settings.HEALTH_CHECK_SECONDS)) # return data @celery_app.task(name="app.controllers.memory_storage_controller.search_all") def write_total_memory_task(workspace_id: str) -> Dict[str, Any]: """定时任务：查询工作空间下所有宿主的记忆总量并写入数据库 Args: workspace_id: 工作空间ID Returns: 包含任务执行结果的字典 """ start_time = time.time() async def _run() -> Dict[str, Any]: from app.models.app_model import App from app.models.end_user_model import EndUser from app.repositories.memory_increment_repository import write_memory_increment from app.services.memory_storage_service import search_all with get_db_context() as db: try: workspace_uuid = uuid.UUID(workspace_id) # 1. 查询当前workspace下的所有app（仅未删除的） apps = db.query(App).filter( App.workspace_id == workspace_uuid, App.is_active.is_(True) ).all() if not apps: # 如果没有app，总量为0 memory_increment = write_memory_increment( db=db, workspace_id=workspace_uuid, total_num=0 ) return { "status": "SUCCESS", "workspace_id": workspace_id, "total_num": 0, "end_user_count": 0, "memory_increment_id": str(memory_increment.id), "created_at": memory_increment.created_at.isoformat(), } # 2. 查询所有app下的end_user_id（去重） app_ids = [app.id for app in apps] end_users = db.query(EndUser.id).filter( EndUser.app_id.in_(app_ids) ).distinct().all() # 3. 遍历所有end_user，查询每个宿主的记忆总量并累加 total_num = 0 end_user_details = [] for (end_user_id,) in end_users: try: # 调用 search_all 接口查询该宿主的总量 result = await search_all(str(end_user_id)) user_total = result.get("total", 0) total_num += user_total end_user_details.append({ "end_user_id": str(end_user_id), "total": user_total }) except Exception as e: # 记录单个用户查询失败，但继续处理其他用户 end_user_details.append({ "end_user_id": str(end_user_id), "total": 0, "error": str(e) }) # 4. 写入数据库 memory_increment = write_memory_increment( db=db, workspace_id=workspace_uuid, total_num=total_num ) return { "status": "SUCCESS", "workspace_id": workspace_id, "total_num": total_num, "end_user_count": len(end_users), "end_user_details": end_user_details, "memory_increment_id": str(memory_increment.id), "created_at": memory_increment.created_at.isoformat(), } except Exception as e: raise e try: result = asyncio.run(_run()) elapsed_time = time.time() - start_time result["elapsed_time"] = elapsed_time return result except Exception as e: elapsed_time = time.time() - start_time return { "status": "FAILURE", "error": str(e), "workspace_id": workspace_id, "elapsed_time": elapsed_time, } @celery_app.task( name="app.tasks.regenerate_memory_cache", bind=True, ignore_result=True, max_retries=0, acks_late=False, time_limit=3600, soft_time_limit=3300, ) def regenerate_memory_cache(self) -> Dict[str, Any]: """定时任务：为所有用户重新生成记忆洞察和用户摘要缓存 遍历所有活动工作空间的所有终端用户，为每个用户重新生成记忆洞察和用户摘要。 实现错误隔离，单个用户失败不影响其他用户的处理。 Returns: 包含任务执行结果的字典，包括： - status: 任务状态 (SUCCESS/FAILURE) - message: 执行消息 - workspace_count: 处理的工作空间数量 - total_users: 总用户数 - successful: 成功生成的用户数 - failed: 失败的用户数 - workspace_results: 每个工作空间的详细结果 - elapsed_time: 执行耗时（秒） - task_id: 任务ID """ start_time = time.time() async def _run() -> Dict[str, Any]: from app.core.logging_config import get_logger from app.repositories.end_user_repository import EndUserRepository from app.services.user_memory_service import UserMemoryService logger = get_logger(__name__) logger.info("开始执行记忆缓存重新生成定时任务") service = UserMemoryService() total_users = 0 successful = 0 failed = 0 workspace_results = [] with get_db_context() as db: try: # 获取所有活动工作空间 repo = EndUserRepository(db) workspaces = repo.get_all_active_workspaces() logger.info(f"找到 {len(workspaces)} 个活动工作空间") # 遍历每个工作空间 for workspace_id in workspaces: logger.info(f"开始处理工作空间: {workspace_id}") workspace_start_time = time.time() try: # 获取工作空间的所有终端用户 end_users = repo.get_all_by_workspace(workspace_id) workspace_user_count = len(end_users) total_users += workspace_user_count logger.info(f"工作空间 {workspace_id} 有 {workspace_user_count} 个终端用户") workspace_successful = 0 workspace_failed = 0 workspace_errors = [] # 遍历每个用户并生成缓存 for end_user in end_users: end_user_id = str(end_user.id) try: # 生成记忆洞察 insight_result = await service.generate_and_cache_insight(db, end_user_id) # 生成用户摘要 summary_result = await service.generate_and_cache_summary(db, end_user_id) # 检查是否都成功 if insight_result["success"] and summary_result["success"]: workspace_successful += 1 successful += 1 logger.info(f"成功为终端用户 {end_user_id} 重新生成缓存") else: workspace_failed += 1 failed += 1 error_info = { "end_user_id": end_user_id, "insight_error": insight_result.get("error"), "summary_error": summary_result.get("error") } workspace_errors.append(error_info) logger.warning(f"终端用户 {end_user_id} 的缓存重新生成部分失败: {error_info}") except Exception as e: # 单个用户失败不影响其他用户（错误隔离） workspace_failed += 1 failed += 1 error_info = { "end_user_id": end_user_id, "error": str(e) } workspace_errors.append(error_info) logger.error(f"为终端用户 {end_user_id} 重新生成缓存时出错: {str(e)}") workspace_elapsed = time.time() - workspace_start_time # 记录工作空间处理结果 workspace_result = { "workspace_id": str(workspace_id), "total_users": workspace_user_count, "successful": workspace_successful, "failed": workspace_failed, "errors": workspace_errors[:10], # 只保留前10个错误 "elapsed_time": workspace_elapsed } workspace_results.append(workspace_result) logger.info( f"工作空间 {workspace_id} 处理完成: " f"总数={workspace_user_count}, 成功={workspace_successful}, " f"失败={workspace_failed}, 耗时={workspace_elapsed:.2f}秒" ) except Exception as e: # 工作空间处理失败，记录错误并继续处理下一个 logger.error(f"处理工作空间 {workspace_id} 时出错: {str(e)}") workspace_results.append({ "workspace_id": str(workspace_id), "error": str(e), "total_users": 0, "successful": 0, "failed": 0, "errors": [] }) # 记录总体统计信息 logger.info( f"记忆缓存重新生成定时任务完成: " f"工作空间数={len(workspaces)}, 总用户数={total_users}, " f"成功={successful}, 失败={failed}" ) return { "status": "SUCCESS", "message": f"成功处理 {len(workspaces)} 个工作空间，总共 {successful}/{total_users} 个用户缓存重新生成成功", "workspace_count": len(workspaces), "total_users": total_users, "successful": successful, "failed": failed, "workspace_results": workspace_results } except Exception as e: logger.error(f"记忆缓存重新生成定时任务执行失败: {str(e)}") return { "status": "FAILURE", "error": str(e), "workspace_count": len(workspace_results), "total_users": total_users, "successful": successful, "failed": failed, "workspace_results": workspace_results } try: # 使用 nest_asyncio 来避免事件循环冲突 try: import nest_asyncio nest_asyncio.apply() except ImportError: pass # 尝试获取现有事件循环，如果不存在则创建新的 try: loop = asyncio.get_event_loop() if loop.is_closed(): loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) except RuntimeError: loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) result = loop.run_until_complete(_run()) elapsed_time = time.time() - start_time result["elapsed_time"] = elapsed_time result["task_id"] = self.request.id return result except Exception as e: elapsed_time = time.time() - start_time return { "status": "FAILURE", "error": str(e), "elapsed_time": elapsed_time, "task_id": self.request.id } @celery_app.task( name="app.tasks.workspace_reflection_task", bind=True, ignore_result=True, max_retries=0, acks_late=False, time_limit=300, soft_time_limit=240, ) def workspace_reflection_task(self) -> Dict[str, Any]: """定时任务：每30秒运行工作空间反思功能 Returns: 包含任务执行结果的字典 """ start_time = time.time() async def _run() -> Dict[str, Any]: from app.core.logging_config import get_api_logger from app.models.workspace_model import Workspace from app.services.memory_reflection_service import ( MemoryReflectionService, WorkspaceAppService, ) api_logger = get_api_logger() with get_db_context() as db: try: # 获取所有工作空间 workspaces = db.query(Workspace).all() if not workspaces: return { "status": "SUCCESS", "message": "没有找到工作空间", "workspace_count": 0, "reflection_results": [] } all_reflection_results = [] # 遍历每个工作空间 for workspace in workspaces: workspace_id = workspace.id api_logger.info(f"开始处理工作空间反思，workspace_id: {workspace_id}") try: reflection_service = MemoryReflectionService(db) # 使用服务类处理复杂查询逻辑 service = WorkspaceAppService(db) result = service.get_workspace_apps_detailed(str(workspace_id)) workspace_reflection_results = [] for data in result['apps_detailed_info']: if data['data_configs'] == []: continue releases = data['releases'] data_configs = data['data_configs'] end_users = data['end_users'] for base, config, user in zip(releases, data_configs, end_users): if str(base['config']) == str(config['config_id']) and str(base['app_id']) == str(user['app_id']): # 调用反思服务 api_logger.info(f"为用户 {user['id']} 启动反思，config_id: {config['config_id']}") reflection_result = await reflection_service.start_reflection_from_data( config_data=config, end_user_id=user['id'] ) workspace_reflection_results.append({ "app_id": base['app_id'], "config_id": config['config_id'], "end_user_id": user['id'], "reflection_result": reflection_result }) all_reflection_results.append({ "workspace_id": str(workspace_id), "reflection_count": len(workspace_reflection_results), "reflection_results": workspace_reflection_results }) api_logger.info( f"工作空间 {workspace_id} 反思处理完成，处理了 {len(workspace_reflection_results)} 个任务") except Exception as e: db.rollback() # Rollback failed transaction to allow next query api_logger.error(f"处理工作空间 {workspace_id} 反思失败: {str(e)}") all_reflection_results.append({ "workspace_id": str(workspace_id), "error": str(e), "reflection_count": 0, "reflection_results": [] }) total_reflections = sum(r.get("reflection_count", 0) for r in all_reflection_results) return { "status": "SUCCESS", "message": f"成功处理 {len(workspaces)} 个工作空间，总共 {total_reflections} 个反思任务", "workspace_count": len(workspaces), "total_reflections": total_reflections, "workspace_results": all_reflection_results } except Exception as e: api_logger.error(f"工作空间反思任务执行失败: {str(e)}") return { "status": "FAILURE", "error": str(e), "workspace_count": 0, "reflection_results": [] } try: # 使用 nest_asyncio 来避免事件循环冲突 try: import nest_asyncio nest_asyncio.apply() except ImportError: pass # 尝试获取现有事件循环，如果不存在则创建新的 try: loop = asyncio.get_event_loop() if loop.is_closed(): loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) except RuntimeError: loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) result = loop.run_until_complete(_run()) elapsed_time = time.time() - start_time result["elapsed_time"] = elapsed_time result["task_id"] = self.request.id return result except Exception as e: elapsed_time = time.time() - start_time return { "status": "FAILURE", "error": str(e), "elapsed_time": elapsed_time, "task_id": self.request.id } @celery_app.task( name="app.tasks.run_forgetting_cycle_task", bind=True, ignore_result=True, max_retries=0, acks_late=False, time_limit=7200, soft_time_limit=7000, ) def run_forgetting_cycle_task(self, config_id: Optional[uuid.UUID] = None) -> Dict[str, Any]: """定时任务：运行遗忘周期 定期执行遗忘周期，识别并融合低激活值的知识节点。 Args: config_id: 配置ID（可选，如果为None则使用默认配置） Returns: 包含任务执行结果的字典 """ start_time = time.time() async def _run() -> Dict[str, Any]: from app.core.logging_config import get_api_logger from app.services.memory_forget_service import MemoryForgetService api_logger = get_api_logger() with get_db_context() as db: try: api_logger.info(f"开始执行遗忘周期定时任务，config_id: {config_id}") forget_service = MemoryForgetService() # 运行遗忘周期 report = await forget_service.trigger_forgetting( db=db, end_user_id=None, # 处理所有组 config_id=config_id ) duration = time.time() - start_time api_logger.info( f"遗忘周期定时任务完成: " f"融合 {report['merged_count']} 对节点, " f"失败 {report['failed_count']} 对, " f"耗时 {duration:.2f} 秒" ) return { "status": "SUCCESS", "message": "遗忘周期执行成功", "report": report, "duration_seconds": duration } except Exception as e: duration = time.time() - start_time api_logger.error(f"遗忘周期定时任务失败: {str(e)}", exc_info=True) return { "status": "FAILED", "message": f"遗忘周期执行失败: {str(e)}", "duration_seconds": duration } # 运行异步函数 loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) try: result = loop.run_until_complete(_run()) return result finally: loop.close() # ============================================================================= # Long-term Memory Storage Tasks (Batched Write Strategies) # ============================================================================= @celery_app.task(name="app.core.memory.agent.long_term_storage.window", bind=True) def long_term_storage_window_task( self, end_user_id: str, langchain_messages: List[Dict[str, Any]], config_id: str, scope: int = 6 ) -> Dict[str, Any]: """Celery task for window-based long-term memory storage. Accumulates messages in Redis buffer until window size (scope) is reached, then writes batched messages to Neo4j. Args: end_user_id: End user identifier langchain_messages: List of messages [{"role": "user/assistant", "content": "..."}] config_id: Memory configuration ID scope: Window size (number of messages before triggering write) Returns: Dict containing task status and metadata """ from app.core.logging_config import get_logger logger = get_logger(__name__) logger.info(f"[LONG_TERM_WINDOW] Starting task - end_user_id={end_user_id}, scope={scope}") start_time = time.time() async def _run() -> Dict[str, Any]: from app.core.memory.agent.langgraph_graph.routing.write_router import window_dialogue from app.core.memory.agent.langgraph_graph.tools.write_tool import chat_data_format from app.core.memory.agent.utils.redis_tool import write_store from app.services.memory_config_service import MemoryConfigService db = next(get_db()) try: # Save to Redis buffer first write_store.save_session_write(end_user_id, await chat_data_format(langchain_messages)) # Load memory config config_service = MemoryConfigService(db) memory_config = config_service.load_memory_config( config_id=config_id, service_name="LongTermStorageTask" ) # Execute window-based dialogue storage await window_dialogue(end_user_id, langchain_messages, memory_config, scope) return {"status": "SUCCESS", "strategy": "window", "scope": scope} finally: db.close() try: import nest_asyncio nest_asyncio.apply() except ImportError: pass try: loop = asyncio.get_event_loop() if loop.is_closed(): loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) except RuntimeError: loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) try: result = loop.run_until_complete(_run()) elapsed_time = time.time() - start_time logger.info(f"[LONG_TERM_WINDOW] Task completed - elapsed_time={elapsed_time:.2f}s") return { **result, "end_user_id": end_user_id, "config_id": config_id, "elapsed_time": elapsed_time, "task_id": self.request.id } except Exception as e: elapsed_time = time.time() - start_time logger.error(f"[LONG_TERM_WINDOW] Task failed - error={str(e)}", exc_info=True) return { "status": "FAILURE", "strategy": "window", "error": str(e), "end_user_id": end_user_id, "config_id": config_id, "elapsed_time": elapsed_time, "task_id": self.request.id } # @celery_app.task(name="app.core.memory.agent.long_term_storage.time", bind=True) # def long_term_storage_time_task( # self, # end_user_id: str, # config_id: str, # time_window: int = 5 # ) -> Dict[str, Any]: # """Celery task for time-based long-term memory storage. # Retrieves recent sessions from Redis within time window and writes to Neo4j. # Args: # end_user_id: End user identifier # config_id: Memory configuration ID # time_window: Time window in minutes for retrieving recent sessions # Returns: # Dict containing task status and metadata # """ # from app.core.logging_config import get_logger # logger = get_logger(__name__) # logger.info(f"[LONG_TERM_TIME] Starting task - end_user_id={end_user_id}, time_window={time_window}") # start_time = time.time() # async def _run() -> Dict[str, Any]: # from app.core.memory.agent.langgraph_graph.routing.write_router import memory_long_term_storage # from app.services.memory_config_service import MemoryConfigService # db = next(get_db()) # try: # # Load memory config # config_service = MemoryConfigService(db) # memory_config = config_service.load_memory_config( # config_id=config_id, # service_name="LongTermStorageTask" # ) # # Execute time-based storage # await memory_long_term_storage(end_user_id, memory_config, time_window) # return {"status": "SUCCESS", "strategy": "time", "time_window": time_window} # finally: # db.close() # try: # import nest_asyncio # nest_asyncio.apply() # except ImportError: # pass # try: # loop = asyncio.get_event_loop() # if loop.is_closed(): # loop = asyncio.new_event_loop() # asyncio.set_event_loop(loop) # except RuntimeError: # loop = asyncio.new_event_loop() # asyncio.set_event_loop(loop) # try: # result = loop.run_until_complete(_run()) # elapsed_time = time.time() - start_time # logger.info(f"[LONG_TERM_TIME] Task completed - elapsed_time={elapsed_time:.2f}s") # return { # **result, # "end_user_id": end_user_id, # "config_id": config_id, # "elapsed_time": elapsed_time, # "task_id": self.request.id # } # except Exception as e: # elapsed_time = time.time() - start_time # logger.error(f"[LONG_TERM_TIME] Task failed - error={str(e)}", exc_info=True) # return { # "status": "FAILURE", # "strategy": "time", # "error": str(e), # "end_user_id": end_user_id, # "config_id": config_id, # "elapsed_time": elapsed_time, # "task_id": self.request.id # } # @celery_app.task(name="app.core.memory.agent.long_term_storage.aggregate", bind=True) # def long_term_storage_aggregate_task( # self, # end_user_id: str, # langchain_messages: List[Dict[str, Any]], # config_id: str # ) -> Dict[str, Any]: # """Celery task for aggregate-based long-term memory storage. # Uses LLM to determine if new messages describe the same event as history. # Only writes to Neo4j if messages represent new information (not duplicates). # Args: # end_user_id: End user identifier # langchain_messages: List of messages [{"role": "user/assistant", "content": "..."}] # config_id: Memory configuration ID # Returns: # Dict containing task status, is_same_event flag, and metadata # """ # from app.core.logging_config import get_logger # logger = get_logger(__name__) # logger.info(f"[LONG_TERM_AGGREGATE] Starting task - end_user_id={end_user_id}") # start_time = time.time() # async def _run() -> Dict[str, Any]: # from app.core.memory.agent.langgraph_graph.routing.write_router import aggregate_judgment # from app.core.memory.agent.langgraph_graph.tools.write_tool import chat_data_format # from app.core.memory.agent.utils.redis_tool import write_store # from app.services.memory_config_service import MemoryConfigService # db = next(get_db()) # try: # # Save to Redis buffer first # write_store.save_session_write(end_user_id, await chat_data_format(langchain_messages)) # # Load memory config # config_service = MemoryConfigService(db) # memory_config = config_service.load_memory_config( # config_id=config_id, # service_name="LongTermStorageTask" # ) # # Execute aggregate judgment # result = await aggregate_judgment(end_user_id, langchain_messages, memory_config) # return { # "status": "SUCCESS", # "strategy": "aggregate", # "is_same_event": result.get("is_same_event", False), # "wrote_to_neo4j": not result.get("is_same_event", False) # } # finally: # db.close() # try: # import nest_asyncio # nest_asyncio.apply() # except ImportError: # pass # try: # loop = asyncio.get_event_loop() # if loop.is_closed(): # loop = asyncio.new_event_loop() # asyncio.set_event_loop(loop) # except RuntimeError: # loop = asyncio.new_event_loop() # asyncio.set_event_loop(loop) # try: # result = loop.run_until_complete(_run()) # elapsed_time = time.time() - start_time # logger.info(f"[LONG_TERM_AGGREGATE] Task completed - is_same_event={result.get('is_same_event')}, elapsed_time={elapsed_time:.2f}s") # return { # **result, # "end_user_id": end_user_id, # "config_id": config_id, # "elapsed_time": elapsed_time, # "task_id": self.request.id # } # except Exception as e: # elapsed_time = time.time() - start_time # logger.error(f"[LONG_TERM_AGGREGATE] Task failed - error={str(e)}", exc_info=True) # return { # "status": "FAILURE", # "strategy": "aggregate", # "error": str(e), # "end_user_id": end_user_id, # "config_id": config_id, # "elapsed_time": elapsed_time, # "task_id": self.request.id # }