♻️ Backend code cleanup and import organization (database/) #1037

Author: WMC001

backend/agents/agent_run_manager.py

@@ -21,23 +21,27 @@ def __new__(cls):
 
     def __init__(self):
         if not self._initialized:
-            self.agent_runs: Dict[int, AgentRunInfo] = {}  # conversation_id -> agent_run_info
+            # conversation_id -> agent_run_info
+            self.agent_runs: Dict[int, AgentRunInfo] = {}
             self._initialized = True
 
     def register_agent_run(self, conversation_id: int, agent_run_info):
         """register agent run instance"""
         with self._lock:
             self.agent_runs[conversation_id] = agent_run_info
-            logger.info(f"register agent run instance, conversation_id: {conversation_id}")
+            logger.info(
+                f"register agent run instance, conversation_id: {conversation_id}")
 
     def unregister_agent_run(self, conversation_id: int):
         """unregister agent run instance"""
         with self._lock:
             if conversation_id in self.agent_runs:
                 del self.agent_runs[conversation_id]
-                logger.info(f"unregister agent run instance, conversation_id: {conversation_id}")
+                logger.info(
+                    f"unregister agent run instance, conversation_id: {conversation_id}")
             else:
-                logger.info(f"no agent run instance found for conversation_id: {conversation_id}")
+                logger.info(
+                    f"no agent run instance found for conversation_id: {conversation_id}")
 
     def get_agent_run_info(self, conversation_id: int):
         """get agent run instance"""
@@ -48,7 +52,8 @@ def stop_agent_run(self, conversation_id: int) -> bool:
         agent_run_info = self.get_agent_run_info(conversation_id)
         if agent_run_info is not None:
             agent_run_info.stop_event.set()
-            logger.info(f"agent run stopped, conversation_id: {conversation_id}")
+            logger.info(
+                f"agent run stopped, conversation_id: {conversation_id}")
             return True
         return False
 

backend/agents/create_agent_info.py

@@ -24,8 +24,10 @@
 
 
 async def create_model_config_list(tenant_id):
-    main_model_config = tenant_config_manager.get_model_config(key="LLM_ID", tenant_id=tenant_id)
-    sub_model_config = tenant_config_manager.get_model_config(key="LLM_SECONDARY_ID", tenant_id=tenant_id)
+    main_model_config = tenant_config_manager.get_model_config(
+        key="LLM_ID", tenant_id=tenant_id)
+    sub_model_config = tenant_config_manager.get_model_config(
+        key="LLM_SECONDARY_ID", tenant_id=tenant_id)
 
     return [ModelConfig(cite_name="main_model",
                         api_key=main_model_config.get("api_key", ""),
@@ -40,10 +42,12 @@ async def create_model_config_list(tenant_id):
 
 
 async def create_agent_config(agent_id, tenant_id, user_id, language: str = 'zh', last_user_query: str = None):
-    agent_info = search_agent_info_by_agent_id(agent_id=agent_id, tenant_id=tenant_id)
+    agent_info = search_agent_info_by_agent_id(
+        agent_id=agent_id, tenant_id=tenant_id)
 
     # create sub agent
-    sub_agent_id_list = query_sub_agents_id_list(main_agent_id=agent_id, tenant_id=tenant_id)
+    sub_agent_id_list = query_sub_agents_id_list(
+        main_agent_id=agent_id, tenant_id=tenant_id)
     managed_agents = []
     for sub_agent_id in sub_agent_id_list:
         sub_agent_config = await create_agent_config(
@@ -55,19 +59,22 @@ async def create_agent_config(agent_id, tenant_id, user_id, language: str = 'zh'
         managed_agents.append(sub_agent_config)
 
     tool_list = await create_tool_config_list(agent_id, tenant_id, user_id)
-    
+
     # Build system prompt: prioritize segmented fields, fallback to original prompt field if not available
     duty_prompt = agent_info.get("duty_prompt", "")
     constraint_prompt = agent_info.get("constraint_prompt", "")
     few_shots_prompt = agent_info.get("few_shots_prompt", "")
-    
+
     # Get template content
-    prompt_template = get_agent_prompt_template(is_manager=len(managed_agents) > 0, language=language)
+    prompt_template = get_agent_prompt_template(
+        is_manager=len(managed_agents) > 0, language=language)
 
     # Get app information
     default_app_description = 'Nexent 是一个开源智能体SDK和平台' if language == 'zh' else 'Nexent is an open-source agent SDK and platform'
-    app_name = tenant_config_manager.get_app_config('APP_NAME', tenant_id=tenant_id) or "Nexent"
-    app_description = tenant_config_manager.get_app_config('APP_DESCRIPTION', tenant_id=tenant_id) or default_app_description
+    app_name = tenant_config_manager.get_app_config(
+        'APP_NAME', tenant_id=tenant_id) or "Nexent"
+    app_description = tenant_config_manager.get_app_config(
+        'APP_DESCRIPTION', tenant_id=tenant_id) or default_app_description
 
     # Get memory list
     memory_context = build_memory_context(user_id, tenant_id, agent_id)
@@ -84,7 +91,8 @@ async def create_agent_config(agent_id, tenant_id, user_id, language: str = 'zh'
             memory_levels.remove("user_agent")
 
         search_res = await search_memory_in_levels(
-            query_text=last_user_query if last_user_query else agent_info.get("name"),
+            query_text=last_user_query if last_user_query else agent_info.get(
+                "name"),
             memory_config=memory_context.memory_config,
             tenant_id=memory_context.tenant_id,
             user_id=memory_context.user_id,
@@ -100,7 +108,8 @@ async def create_agent_config(agent_id, tenant_id, user_id, language: str = 'zh'
     try:
         for tool in tool_list:
             if "KnowledgeBaseSearchTool" == tool.class_name:
-                knowledge_info_list = get_selected_knowledge_list(tenant_id=tenant_id, user_id=user_id)
+                knowledge_info_list = get_selected_knowledge_list(
+                    tenant_id=tenant_id, user_id=user_id)
                 if knowledge_info_list:
                     for knowledge_info in knowledge_info_list:
                         knowledge_name = knowledge_info.get("index_name")
@@ -109,13 +118,14 @@ async def create_agent_config(agent_id, tenant_id, user_id, language: str = 'zh'
                             summary = message.get("summary", "")
                             knowledge_base_summary += f"**{knowledge_name}**: {summary}\n\n"
                         except Exception as e:
-                            logger.warning(f"Failed to get summary for knowledge base {knowledge_name}: {e}")
+                            logger.warning(
+                                f"Failed to get summary for knowledge base {knowledge_name}: {e}")
                 else:
                     knowledge_base_summary = "当前没有可用的知识库索引。\n" if language == 'zh' else "No knowledge base indexes are currently available.\n"
                 break  # Only process the first KnowledgeBaseSearchTool found
     except Exception as e:
         logger.error(f"Failed to build knowledge base summary: {e}")
-    
+
     # Assemble system_prompt
     if duty_prompt or constraint_prompt or few_shots_prompt:
         system_prompt = Template(prompt_template["system_prompt"], undefined=StrictUndefined).render({
@@ -133,7 +143,7 @@ async def create_agent_config(agent_id, tenant_id, user_id, language: str = 'zh'
         })
     else:
         system_prompt = agent_info.get("prompt", "")
-    
+
     agent_config = AgentConfig(
         name="undefined" if agent_info["name"] is None else agent_info["name"],
         description="undefined" if agent_info["description"] is None else agent_info["description"],
@@ -182,25 +192,27 @@ async def create_tool_config_list(agent_id, tenant_id, user_id):
 
         # special logic for knowledge base search tool
         if tool_config.class_name == "KnowledgeBaseSearchTool":
-            knowledge_info_list = get_selected_knowledge_list(tenant_id=tenant_id, user_id=user_id)
-            index_names = [knowledge_info.get("index_name") for knowledge_info in knowledge_info_list]
+            knowledge_info_list = get_selected_knowledge_list(
+                tenant_id=tenant_id, user_id=user_id)
+            index_names = [knowledge_info.get(
+                "index_name") for knowledge_info in knowledge_info_list]
             tool_config.metadata = {"index_names": index_names,
                                     "es_core": elastic_core,
                                     "embedding_model": get_embedding_model(tenant_id=tenant_id)}
         tool_config_list.append(tool_config)
-    
+
     return tool_config_list
 
 
 async def discover_langchain_tools():
     """
     Discover LangChain tools implemented with the `@tool` decorator.
-    
+
     Returns:
         list: List of discovered LangChain tool instances
     """
     from utils.langchain_utils import discover_langchain_modules
-    
+
     langchain_tools = []
 
     # ----------------------------------------------
@@ -210,18 +222,21 @@ async def discover_langchain_tools():
     try:
         # Use the utility function to discover all BaseTool objects
         discovered_tools = discover_langchain_modules()
-        
+
         for obj, filename in discovered_tools:
             try:
                 # Log successful tool discovery
-                logger.info(f"Loaded LangChain tool '{obj.name}' from {filename}")
+                logger.info(
+                    f"Loaded LangChain tool '{obj.name}' from {filename}")
                 langchain_tools.append(obj)
             except Exception as e:
-                logger.error(f"Error processing LangChain tool from {filename}: {e}")
-                
+                logger.error(
+                    f"Error processing LangChain tool from {filename}: {e}")
+
     except Exception as e:
-        logger.error(f"Unexpected error scanning LangChain tools directory: {e}")
-    
+        logger.error(
+            f"Unexpected error scanning LangChain tools directory: {e}")
+
     return langchain_tools
 
 
@@ -269,7 +284,8 @@ def check_agent_tools(agent_config: AgentConfig):
         # Check current agent tools
         for tool in agent_config.tools:
             if tool.source == "mcp" and tool.usage in mcp_info_dict:
-                used_mcp_urls.add(mcp_info_dict[tool.usage]["remote_mcp_server"])
+                used_mcp_urls.add(
+                    mcp_info_dict[tool.usage]["remote_mcp_server"])
 
         # Recursively check sub-agent
         for sub_agent_config in agent_config.managed_agents:
@@ -295,13 +311,15 @@ async def create_agent_run_info(agent_id, minio_files, query, history, authoriza
     )
 
     remote_mcp_list = await get_remote_mcp_server_list(tenant_id=tenant_id)
-    default_mcp_url = urljoin(config_manager.get_config("NEXENT_MCP_SERVER"), "sse")
+    default_mcp_url = urljoin(
+        config_manager.get_config("NEXENT_MCP_SERVER"), "sse")
     remote_mcp_list.append({
         "remote_mcp_server_name": "nexent",
         "remote_mcp_server": default_mcp_url,
         "status": True
     })
-    remote_mcp_dict = {record["remote_mcp_server_name"]: record for record in remote_mcp_list if record["status"]}
+    remote_mcp_dict = {record["remote_mcp_server_name"]
+        : record for record in remote_mcp_list if record["status"]}
 
     # Filter MCP servers and tools
     mcp_host = filter_mcp_servers_and_tools(agent_config, remote_mcp_dict)

backend/agents/preprocess_manager.py

@@ -30,8 +30,10 @@ def __new__(cls):
 
     def __init__(self):
         if not self._initialized:
-            self.preprocess_tasks: Dict[str, PreprocessTask] = {}  # task_id -> PreprocessTask
-            self.conversation_tasks: Dict[int, Set[str]] = {}  # conversation_id -> Set[task_id]
+            # task_id -> PreprocessTask
+            self.preprocess_tasks: Dict[str, PreprocessTask] = {}
+            # conversation_id -> Set[task_id]
+            self.conversation_tasks: Dict[int, Set[str]] = {}
             self._initialized = True
 
     def register_preprocess_task(self, task_id: str, conversation_id: int, task: asyncio.Task):
@@ -40,76 +42,78 @@ def register_preprocess_task(self, task_id: str, conversation_id: int, task: asy
             preprocess_task = PreprocessTask(task_id, conversation_id)
             preprocess_task.task = task
             self.preprocess_tasks[task_id] = preprocess_task
-            
+
             if conversation_id not in self.conversation_tasks:
                 self.conversation_tasks[conversation_id] = set()
             self.conversation_tasks[conversation_id].add(task_id)
-            
-            logger.info(f"Registered preprocess task {task_id} for conversation {conversation_id}")
+
+            logger.info(
+                f"Registered preprocess task {task_id} for conversation {conversation_id}")
 
     def unregister_preprocess_task(self, task_id: str):
         """Unregister a preprocess task"""
         with self._lock:
             if task_id in self.preprocess_tasks:
                 task = self.preprocess_tasks[task_id]
                 conversation_id = task.conversation_id
-                
+
                 # Remove from conversation_tasks
                 if conversation_id in self.conversation_tasks:
                     self.conversation_tasks[conversation_id].discard(task_id)
                     if not self.conversation_tasks[conversation_id]:
                         del self.conversation_tasks[conversation_id]
-                
+
                 # Remove from preprocess_tasks
                 del self.preprocess_tasks[task_id]
-                
+
                 logger.info(f"Unregistered preprocess task {task_id}")
 
     def stop_preprocess_tasks(self, conversation_id: int) -> bool:
         """Stop all preprocess tasks for a conversation"""
         with self._lock:
             if conversation_id not in self.conversation_tasks:
                 return False
-            
+
             task_ids = self.conversation_tasks[conversation_id].copy()
             stopped_count = 0
-            
+
             for task_id in task_ids:
                 if task_id in self.preprocess_tasks:
                     task = self.preprocess_tasks[task_id]
                     if task.is_running:
                         task.stop_event.set()
                         task.is_running = False
-                        
+
                         # Cancel the asyncio task if it exists
                         if task.task and not task.task.done():
                             task.task.cancel()
-                        
+
                         stopped_count += 1
-                        logger.info(f"Stopped preprocess task {task_id} for conversation {conversation_id}")
-            
+                        logger.info(
+                            f"Stopped preprocess task {task_id} for conversation {conversation_id}")
+
             return stopped_count > 0
 
     def is_preprocess_running(self, conversation_id: int) -> bool:
         """Check if any preprocess task is running for a conversation"""
         with self._lock:
             if conversation_id not in self.conversation_tasks:
                 return False
-            
+
             for task_id in self.conversation_tasks[conversation_id]:
                 if task_id in self.preprocess_tasks:
                     task = self.preprocess_tasks[task_id]
                     if task.is_running and not task.stop_event.is_set():
                         return True
-            
+
             return False
 
     def get_preprocess_status(self, conversation_id: int) -> Dict:
         """Get preprocess status for a conversation"""
         with self._lock:
             if conversation_id not in self.conversation_tasks:
                 return {"running": False, "task_count": 0}
-            
+
             running_tasks = []
             for task_id in self.conversation_tasks[conversation_id]:
                 if task_id in self.preprocess_tasks:
@@ -119,7 +123,7 @@ def get_preprocess_status(self, conversation_id: int) -> Dict:
                         "is_running": task.is_running,
                         "stopped": task.stop_event.is_set()
                     })
-            
+
             return {
                 "running": any(task["is_running"] for task in running_tasks),
                 "task_count": len(running_tasks),

backend/apps/agent_app.py

@@ -139,7 +139,6 @@ async def import_agent_api(request: AgentImportRequest, authorization: Optional[
             status_code=HTTPStatus.INTERNAL_SERVER_ERROR, detail="Agent import error.")
 
 
-
 @router.get("/list")
 async def list_all_agent_info_api(authorization: Optional[str] = Header(None), request: Request = None):
     """