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

fix(db): fix database connection handling

Browse Source
This commit is contained in:
mengyonghao
2025-12-24 12:21:12 +08:00
parent b99671e04a
commit 38220006a6
1 changed files with 60 additions and 64 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
api/app/core/workflow/nodes/knowledge/node.py
View File

@@ -20,9 +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_gen = get_db() db = next(get_db())
db = next(db_gen)
try:
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,
@@ -93,5 +91,3 @@ class KnowledgeRetrievalNode(BaseNode):
unique_rs.append(doc) unique_rs.append(doc)
rs = vector_service.rerank(query=query, docs=unique_rs, top_k=self.typed_config.top_k) rs = vector_service.rerank(query=query, docs=unique_rs, top_k=self.typed_config.top_k)
return [chunk.model_dump() for chunk in rs] return [chunk.model_dump() for chunk in rs]
finally:
next(db_gen)