fix(db): fix database connection handling

This commit is contained in:
mengyonghao
2025-12-24 12:22:59 +08:00
parent 38220006a6
commit 0a8c1be084

View File

@@ -4,7 +4,7 @@ from typing import Any
from app.core.rag.vdb.elasticsearch.elasticsearch_vector import ElasticSearchVectorFactory from app.core.rag.vdb.elasticsearch.elasticsearch_vector import ElasticSearchVectorFactory
from app.core.workflow.nodes.base_node import BaseNode, WorkflowState from app.core.workflow.nodes.base_node import BaseNode, WorkflowState
from app.core.workflow.nodes.knowledge import KnowledgeRetrievalNodeConfig from app.core.workflow.nodes.knowledge import KnowledgeRetrievalNodeConfig
from app.db import get_db from app.db import get_db_context
from app.models import knowledge_model, knowledgeshare_model from app.models import knowledge_model, knowledgeshare_model
from app.repositories import knowledge_repository from app.repositories import knowledge_repository
from app.schemas.chunk_schema import RetrieveType from app.schemas.chunk_schema import RetrieveType
@@ -20,7 +20,7 @@ class KnowledgeRetrievalNode(BaseNode):
async def execute(self, state: WorkflowState) -> Any: async def execute(self, state: WorkflowState) -> Any:
query = self._render_template(self.typed_config.query, state) query = self._render_template(self.typed_config.query, state)
db = next(get_db()) with get_db_context():
filters = [ filters = [
knowledge_model.Knowledge.id.in_(self.typed_config.kb_ids), knowledge_model.Knowledge.id.in_(self.typed_config.kb_ids),
knowledge_model.Knowledge.permission_id == knowledge_model.PermissionType.Private, knowledge_model.Knowledge.permission_id == knowledge_model.PermissionType.Private,