fix(db): fix database connection leak

api/app/controllers/memory_agent_controller.py

@@ -1,28 +1,29 @@
 from typing import List, Optional
 
+from dotenv import load_dotenv
+from fastapi import APIRouter, Depends, File, Form, Query, UploadFile, Header
+from sqlalchemy.orm import Session
+from starlette.responses import StreamingResponse
+
 from app.cache.memory.interest_memory import InterestMemoryCache
 from app.celery_app import celery_app
 from app.core.error_codes import BizCode
 from app.core.language_utils import get_language_from_header
 from app.core.logging_config import get_api_logger
+from app.core.memory.agent.utils.redis_tool import store
+from app.core.memory.agent.utils.session_tools import SessionService
 from app.core.rag.llm.cv_model import QWenCV
 from app.core.response_utils import fail, success
 from app.db import get_db
 from app.dependencies import cur_workspace_access_guard, get_current_user
 from app.models import ModelApiKey
 from app.models.user_model import User
-from app.core.memory.agent.utils.session_tools import SessionService
-from app.core.memory.agent.utils.redis_tool import store
-from app.repositories import knowledge_repository, WorkspaceRepository
+from app.repositories import knowledge_repository
 from app.schemas.memory_agent_schema import UserInput, Write_UserInput
 from app.schemas.response_schema import ApiResponse
 from app.services import task_service, workspace_service
 from app.services.memory_agent_service import MemoryAgentService
 from app.services.model_service import ModelConfigService
-from dotenv import load_dotenv
-from fastapi import APIRouter, Depends, File, Form, Query, UploadFile,Header
-from sqlalchemy.orm import Session
-from starlette.responses import StreamingResponse
 
 load_dotenv()
 api_logger = get_api_logger()
@@ -37,7 +38,7 @@ router = APIRouter(
 
 @router.get("/health/status", response_model=ApiResponse)
 async def get_health_status(
-    current_user: User = Depends(get_current_user)
+        current_user: User = Depends(get_current_user)
 ):
     """
     Get latest health status written by Celery periodic task
@@ -55,8 +56,9 @@ async def get_health_status(
 
 @router.get("/download_log")
 async def download_log(
-    log_type: str = Query("file", regex="^(file|transmission)$", description="日志类型: file=完整文件, transmission=实时流式传输"),
-    current_user: User = Depends(get_current_user)
+        log_type: str = Query("file", regex="^(file|transmission)$",
+                              description="日志类型: file=完整文件, transmission=实时流式传输"),
+        current_user: User = Depends(get_current_user)
 ):
     """
     Download or stream agent service log file
@@ -75,16 +77,16 @@ async def download_log(
         - transmission mode: StreamingResponse with SSE
     """
     api_logger.info(f"Log download requested with log_type={log_type}")
-    
+
     # Validate log_type parameter (FastAPI Query regex already validates, but explicit check for clarity)
     if log_type not in ["file", "transmission"]:
         api_logger.warning(f"Invalid log_type parameter: {log_type}")
         return fail(
-            BizCode.BAD_REQUEST, 
-            "无效的log_type参数", 
+            BizCode.BAD_REQUEST,
+            "无效的log_type参数",
             "log_type必须是'file'或'transmission'"
         )
-    
+
     # Route to appropriate mode
     if log_type == "file":
         # File mode: Return complete log file content
@@ -119,10 +121,10 @@ async def download_log(
 @router.post("/writer_service", response_model=ApiResponse)
 @cur_workspace_access_guard()
 async def write_server(
-    user_input: Write_UserInput,
-    language_type: str = Header(default=None, alias="X-Language-Type"),
-    db: Session = Depends(get_db),
-    current_user: User = Depends(get_current_user)
+        user_input: Write_UserInput,
+        language_type: str = Header(default=None, alias="X-Language-Type"),
+        db: Session = Depends(get_db),
+        current_user: User = Depends(get_current_user)
 ):
     """
     Write service endpoint - processes write operations synchronously
@@ -136,11 +138,11 @@ async def write_server(
     """
     # 使用集中化的语言校验
     language = get_language_from_header(language_type)
-    
+
     config_id = user_input.config_id
     workspace_id = current_user.current_workspace_id
     api_logger.info(f"Write service: workspace_id={workspace_id}, config_id={config_id}, language_type={language}")
-    
+
     # 获取 storage_type，如果为 None 则使用默认值
     storage_type = workspace_service.get_workspace_storage_type(
         db=db,
@@ -149,7 +151,7 @@ async def write_server(
     )
     if storage_type is None: storage_type = 'neo4j'
     user_rag_memory_id = ''
-    
+
     # 如果 storage_type 是 rag，必须确保有有效的 user_rag_memory_id
     if storage_type == 'rag':
         if workspace_id:
@@ -161,13 +163,15 @@ async def write_server(
             if knowledge:
                 user_rag_memory_id = str(knowledge.id)
             else:
-                api_logger.warning(f"未找到名为 'USER_RAG_MERORY' 的知识库，workspace_id: {workspace_id}，将使用 neo4j 存储")
+                api_logger.warning(
+                    f"未找到名为 'USER_RAG_MERORY' 的知识库，workspace_id: {workspace_id}，将使用 neo4j 存储")
                 storage_type = 'neo4j'
         else:
             api_logger.warning("workspace_id 为空，无法使用 rag 存储，将使用 neo4j 存储")
             storage_type = 'neo4j'
-    
-    api_logger.info(f"Write service requested for group {user_input.end_user_id}, storage_type: {storage_type}, user_rag_memory_id: {user_rag_memory_id}")
+
+    api_logger.info(
+        f"Write service requested for group {user_input.end_user_id}, storage_type: {storage_type}, user_rag_memory_id: {user_rag_memory_id}")
     try:
         messages_list = memory_agent_service.get_messages_list(user_input)
         result = await memory_agent_service.write_memory(
@@ -175,7 +179,7 @@ async def write_server(
             messages_list,
             config_id,
             db,
-            storage_type, 
+            storage_type,
             user_rag_memory_id,
             language
         )
@@ -195,10 +199,10 @@ async def write_server(
 @router.post("/writer_service_async", response_model=ApiResponse)
 @cur_workspace_access_guard()
 async def write_server_async(
-    user_input: Write_UserInput,
-    language_type: str = Header(default=None, alias="X-Language-Type"),
-    db: Session = Depends(get_db),
-    current_user: User = Depends(get_current_user)
+        user_input: Write_UserInput,
+        language_type: str = Header(default=None, alias="X-Language-Type"),
+        db: Session = Depends(get_db),
+        current_user: User = Depends(get_current_user)
 ):
     """
     Async write service endpoint - enqueues write processing to Celery
@@ -213,10 +217,11 @@ async def write_server_async(
     """
     # 使用集中化的语言校验
     language = get_language_from_header(language_type)
-    
+
     config_id = user_input.config_id
     workspace_id = current_user.current_workspace_id
-    api_logger.info(f"Async write service: workspace_id={workspace_id}, config_id={config_id}, language_type={language}")
+    api_logger.info(
+        f"Async write service: workspace_id={workspace_id}, config_id={config_id}, language_type={language}")
 
     # 获取 storage_type，如果为 None 则使用默认值
     storage_type = workspace_service.get_workspace_storage_type(
@@ -244,7 +249,7 @@ async def write_server_async(
             args=[user_input.end_user_id, messages_list, config_id, storage_type, user_rag_memory_id, language]
         )
         api_logger.info(f"Write task queued: {task.id}")
-        
+
         return success(data={"task_id": task.id}, msg="写入任务已提交")
     except Exception as e:
         api_logger.error(f"Async write operation failed: {str(e)}")
@@ -254,9 +259,9 @@ async def write_server_async(
 @router.post("/read_service", response_model=ApiResponse)
 @cur_workspace_access_guard()
 async def read_server(
-    user_input: UserInput,
-    db: Session = Depends(get_db),
-    current_user: User = Depends(get_current_user)
+        user_input: UserInput,
+        db: Session = Depends(get_db),
+        current_user: User = Depends(get_current_user)
 ):
     """
     Read service endpoint - processes read operations synchronously
@@ -291,8 +296,9 @@ async def read_server(
         )
         if knowledge:
             user_rag_memory_id = str(knowledge.id)
-    
-    api_logger.info(f"Read service: group={user_input.end_user_id}, storage_type={storage_type}, user_rag_memory_id={user_rag_memory_id}, workspace_id={workspace_id}")
+
+    api_logger.info(
+        f"Read service: group={user_input.end_user_id}, storage_type={storage_type}, user_rag_memory_id={user_rag_memory_id}, workspace_id={workspace_id}")
     try:
         result = await memory_agent_service.read_memory(
             user_input.end_user_id,
@@ -306,7 +312,8 @@ async def read_server(
         )
         if str(user_input.search_switch) == "2":
             retrieve_info = result['answer']
-            history = await SessionService(store).get_history(user_input.end_user_id, user_input.end_user_id, user_input.end_user_id)
+            history = await SessionService(store).get_history(user_input.end_user_id, user_input.end_user_id,
+                                                              user_input.end_user_id)
             query = user_input.message
 
             # 调用 memory_agent_service 的方法生成最终答案
@@ -319,7 +326,7 @@ async def read_server(
                 db=db
             )
             if "信息不足，无法回答" in result['answer']:
-                result['answer']=retrieve_info
+                result['answer'] = retrieve_info
         return success(data=result, msg="回复对话消息成功")
     except BaseException as e:
         # Handle ExceptionGroup from TaskGroup (Python 3.11+) or BaseExceptionGroup
@@ -335,9 +342,10 @@ async def read_server(
 @router.post("/file", response_model=ApiResponse)
 async def file_update(
         files: List[UploadFile] = File(..., description="要上传的文件"),
-        model_id:str = Form(..., description="模型ID"),
+        model_id: str = Form(..., description="模型ID"),
         metadata: Optional[str] = Form(None, description="文件元数据 (JSON格式)"),
-        current_user: User = Depends(get_current_user)
+        current_user: User = Depends(get_current_user),
+        db: Session = Depends(get_db),
 ):
     """
     文件上传接口 - 支持图片识别
@@ -350,9 +358,6 @@ async def file_update(
     Returns:
         文件处理结果
     """
-
-    db_gen = get_db()             # get_db 通常是一个生成器
-    db = next(db_gen)
     api_logger.info(f"File upload requested, file count: {len(files)}")
     config = ModelConfigService.get_model_by_id(db=db, model_id=model_id)
     apiConfig: ModelApiKey = config.api_keys[0]
@@ -361,7 +366,7 @@ async def file_update(
         for file in files:
             api_logger.debug(f"Processing file: {file.filename}, content_type: {file.content_type}")
             content = await file.read()
-            
+
             if file.content_type and file.content_type.startswith("image/"):
                 vision_model = QWenCV(
                     key=apiConfig.api_key,
@@ -375,12 +380,12 @@ async def file_update(
             else:
                 api_logger.warning(f"Unsupported file type: {file.content_type}")
                 file_content.append(f"[不支持的文件类型: {file.content_type}]")
-        
+
         result_text = ';'.join(file_content)
         api_logger.info(f"File processing completed, result length: {len(result_text)}")
-        
+
         return success(data=result_text, msg="转换文本成功")
-        
+
     except Exception as e:
         api_logger.error(f"File processing failed: {str(e)}", exc_info=True)
         return fail(BizCode.INTERNAL_ERROR, "转换文本失败", str(e))
@@ -430,8 +435,8 @@ async def read_server_async(
 
 @router.get("/read_result/", response_model=ApiResponse)
 async def get_read_task_result(
-    task_id: str,
-    current_user: User = Depends(get_current_user)
+        task_id: str,
+        current_user: User = Depends(get_current_user)
 ):
     """
     Get the status and result of an async read task
@@ -452,7 +457,7 @@ async def get_read_task_result(
     try:
         result = task_service.get_task_memory_read_result(task_id)
         status = result.get("status")
-        
+
         if status == "SUCCESS":
             # 任务成功完成
             task_result = result.get("result", {})
@@ -470,7 +475,7 @@ async def get_read_task_result(
             else:
                 # 旧格式：直接返回结果
                 return success(data=task_result, msg="查询任务已完成")
-                
+
         elif status == "FAILURE":
             # 任务失败
             error_info = result.get("result", "Unknown error")
@@ -479,7 +484,7 @@ async def get_read_task_result(
             else:
                 error_msg = str(error_info)
             return fail(BizCode.INTERNAL_ERROR, "查询任务失败", error_msg)
-            
+
         elif status in ["PENDING", "STARTED"]:
             # 任务进行中
             return success(
@@ -499,7 +504,7 @@ async def get_read_task_result(
                 },
                 msg=f"任务状态: {status}"
             )
-            
+
     except Exception as e:
         api_logger.error(f"Read task status check failed: {str(e)}", exc_info=True)
         return fail(BizCode.INTERNAL_ERROR, "任务状态查询失败", str(e))
@@ -507,8 +512,8 @@ async def get_read_task_result(
 
 @router.get("/write_result/", response_model=ApiResponse)
 async def get_write_task_result(
-    task_id: str,
-    current_user: User = Depends(get_current_user)
+        task_id: str,
+        current_user: User = Depends(get_current_user)
 ):
     """
     Get the status and result of an async write task
@@ -529,7 +534,7 @@ async def get_write_task_result(
     try:
         result = task_service.get_task_memory_write_result(task_id)
         status = result.get("status")
-        
+
         if status == "SUCCESS":
             # 任务成功完成
             task_result = result.get("result", {})
@@ -547,7 +552,7 @@ async def get_write_task_result(
             else:
                 # 旧格式：直接返回结果
                 return success(data=task_result, msg="写入任务已完成")
-                
+
         elif status == "FAILURE":
             # 任务失败
             error_info = result.get("result", "Unknown error")
@@ -556,7 +561,7 @@ async def get_write_task_result(
             else:
                 error_msg = str(error_info)
             return fail(BizCode.INTERNAL_ERROR, "写入任务失败", error_msg)
-            
+
         elif status in ["PENDING", "STARTED"]:
             # 任务进行中
             return success(
@@ -576,7 +581,7 @@ async def get_write_task_result(
                 },
                 msg=f"任务状态: {status}"
             )
-            
+
     except Exception as e:
         api_logger.error(f"Write task status check failed: {str(e)}", exc_info=True)
         return fail(BizCode.INTERNAL_ERROR, "任务状态查询失败", str(e))
@@ -584,9 +589,9 @@ async def get_write_task_result(
 
 @router.post("/status_type", response_model=ApiResponse)
 async def status_type(
-    user_input: Write_UserInput,
-    db: Session = Depends(get_db),
-    current_user: User = Depends(get_current_user)
+        user_input: Write_UserInput,
+        db: Session = Depends(get_db),
+        current_user: User = Depends(get_current_user)
 ):
     """
     Determine the type of user message (read or write)
@@ -629,9 +634,10 @@ async def status_type(
 
 @router.get("/stats/types", response_model=ApiResponse)
 async def get_knowledge_type_stats_api(
-    end_user_id: Optional[str] = Query(None, description="用户ID（可选）"),
-    only_active: bool = Query(True, description="仅统计有效记录(status=1)"),
-    current_user: User = Depends(get_current_user)
+        end_user_id: Optional[str] = Query(None, description="用户ID（可选）"),
+        only_active: bool = Query(True, description="仅统计有效记录(status=1)"),
+        current_user: User = Depends(get_current_user),
+        db: Session = Depends(get_db),
 ):
     """
     统计当前空间下各知识库类型的数量，包含 General | Web | Third-party | Folder。
@@ -640,14 +646,9 @@ async def get_knowledge_type_stats_api(
     - 知识库类型根据当前用户的 current_workspace_id 过滤
     - 如果用户没有当前工作空间，对应的统计返回 0
     """
-    api_logger.info(f"Knowledge type stats requested for workspace_id: {current_user.current_workspace_id}, end_user_id: {end_user_id}")
+    api_logger.info(
+        f"Knowledge type stats requested for workspace_id: {current_user.current_workspace_id}, end_user_id: {end_user_id}")
     try:
-        from app.db import get_db
-        
-        # 获取数据库会话
-        db_gen = get_db()
-        db = next(db_gen)
-        
         # 调用service层函数
         result = await memory_agent_service.get_knowledge_type_stats(
             end_user_id=end_user_id,
@@ -655,7 +656,7 @@ async def get_knowledge_type_stats_api(
             current_workspace_id=current_user.current_workspace_id,
             db=db
         )
-        
+
         return success(data=result, msg="获取知识库类型统计成功")
     except Exception as e:
         api_logger.error(f"Knowledge type stats failed: {str(e)}")
@@ -664,11 +665,11 @@ async def get_knowledge_type_stats_api(
 
 @router.get("/analytics/interest_distribution/by_user", response_model=ApiResponse)
 async def get_interest_distribution_by_user_api(
-    end_user_id: str = Query(..., description="用户ID（必填）"),
-    limit: int = Query(5, le=5, description="返回兴趣标签数量限制，最多5个"),
-    language_type: str = Header(default=None, alias="X-Language-Type"),
-    current_user: User = Depends(get_current_user),
-    db: Session = Depends(get_db),
+        end_user_id: str = Query(..., description="用户ID（必填）"),
+        limit: int = Query(5, le=5, description="返回兴趣标签数量限制，最多5个"),
+        language_type: str = Header(default=None, alias="X-Language-Type"),
+        current_user: User = Depends(get_current_user),
+        db: Session = Depends(get_db),
 ):
     """
     获取指定用户的兴趣分布标签
@@ -716,9 +717,9 @@ async def get_interest_distribution_by_user_api(
 
 @router.get("/analytics/user_profile", response_model=ApiResponse)
 async def get_user_profile_api(
-    end_user_id: Optional[str] = Query(None, description="用户ID（可选）"),
-    db: Session = Depends(get_db),
-    current_user: User = Depends(get_current_user)
+        end_user_id: Optional[str] = Query(None, description="用户ID（可选）"),
+        db: Session = Depends(get_db),
+        current_user: User = Depends(get_current_user)
 ):
     """
     获取用户详情，包含：
@@ -756,17 +757,17 @@ async def get_user_profile_api(
 # ):
 #     """
 #     Get parsed API documentation (Public endpoint - no authentication required)
-    
+
 #     Args:
 #         file_path: Optional path to API docs file. If None, uses default path.
-    
+
 #     Returns:
 #         Parsed API documentation including title, meta info, and sections
 #     """
 #     api_logger.info(f"API docs requested, file_path: {file_path or 'default'}")
 #     try:
 #         result = await memory_agent_service.get_api_docs(file_path)
-        
+
 #         if result.get("success"):
 #             return success(msg=result["msg"], data=result["data"])
 #         else:
@@ -782,9 +783,9 @@ async def get_user_profile_api(
 
 @router.get("/end_user/{end_user_id}/connected_config", response_model=ApiResponse)
 async def get_end_user_connected_config(
-    end_user_id: str,
-    db: Session = Depends(get_db),
-    current_user: User = Depends(get_current_user)
+        end_user_id: str,
+        db: Session = Depends(get_db),
+        current_user: User = Depends(get_current_user)
 ):
     """
     获取终端用户关联的记忆配置
@@ -803,9 +804,9 @@ async def get_end_user_connected_config(
     from app.services.memory_agent_service import (
         get_end_user_connected_config as get_config,
     )
-    
+
     api_logger.info(f"Getting connected config for end_user: {end_user_id}")
-    
+
     try:
         result = get_config(end_user_id, db)
         return success(data=result, msg="获取终端用户关联配置成功")
@@ -814,4 +815,4 @@ async def get_end_user_connected_config(
         return fail(BizCode.NOT_FOUND, str(e))
     except Exception as e:
         api_logger.error(f"Failed to get end user connected config: {str(e)}", exc_info=True)
-        return fail(BizCode.INTERNAL_ERROR, "获取终端用户关联配置失败", str(e))
+        return fail(BizCode.INTERNAL_ERROR, "获取终端用户关联配置失败", str(e))