outman/dify
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: replace direct ChildChunk queries with SegmentService.get_child_chunk_by_id for improved code clarity and maintainability

Browse Source
This commit is contained in:
ZeroZ_JQ 2025-03-18 18:25:23 +08:00
parent bfe5eb94c0
commit b53888a79e
2 changed files with 17 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
api/controllers/service_api/dataset/segment.py
View File

@ -14,7 +14,7 @@ from core.model_manager import ModelManager
from core.model_runtime.entities.model_entities import ModelType from core.model_runtime.entities.model_entities import ModelType
from extensions.ext_database import db from extensions.ext_database import db
from fields.segment_fields import child_chunk_fields, segment_fields from fields.segment_fields import child_chunk_fields, segment_fields
from models.dataset import ChildChunk, Dataset from models.dataset import Dataset
from services.dataset_service import DatasetService, DocumentService, SegmentService from services.dataset_service import DatasetService, DocumentService, SegmentService
from services.entities.knowledge_entities.knowledge_entities import ChildChunkUpdateArgs, SegmentUpdateArgs from services.entities.knowledge_entities.knowledge_entities import ChildChunkUpdateArgs, SegmentUpdateArgs
from services.errors.chunk import ( from services.errors.chunk import (
@ -336,9 +336,10 @@ class DatasetChildChunkApi(DatasetApiResource):
# check child chunk # check child chunk
child_chunk_id = str(child_chunk_id) child_chunk_id = str(child_chunk_id)
child_chunk = ChildChunk.query.filter( child_chunk = SegmentService.get_child_chunk_by_id(
ChildChunk.id == str(child_chunk_id), ChildChunk.tenant_id == current_user.current_tenant_id child_chunk_id=child_chunk_id,
).first() tenant_id=current_user.current_tenant_id
)
if not child_chunk: if not child_chunk:
raise NotFound("Child chunk not found.") raise NotFound("Child chunk not found.")
@ -376,9 +377,10 @@ class DatasetChildChunkApi(DatasetApiResource):
# check child chunk # check child chunk
child_chunk_id = str(child_chunk_id) child_chunk_id = str(child_chunk_id)
child_chunk = ChildChunk.query.filter( child_chunk = SegmentService.get_child_chunk_by_id(
ChildChunk.id == str(child_chunk_id), ChildChunk.tenant_id == current_user.current_tenant_id child_chunk_id=child_chunk_id,
).first() tenant_id=current_user.current_tenant_id
)
if not child_chunk: if not child_chunk:
raise NotFound("Child chunk not found.") raise NotFound("Child chunk not found.")

8
api/services/dataset_service.py
View File

@ -2155,6 +2155,14 @@ class SegmentService:
query = query.where(ChildChunk.content.ilike(f"%{keyword}%")) query = query.where(ChildChunk.content.ilike(f"%{keyword}%"))
return query.paginate(page=page, per_page=limit, max_per_page=100, error_out=False) return query.paginate(page=page, per_page=limit, max_per_page=100, error_out=False)
@classmethod
def get_child_chunk_by_id(cls, child_chunk_id: str, tenant_id: str) -> Optional[ChildChunk]:
"""Get a child chunk by its ID."""
return ChildChunk.query.filter(
ChildChunk.id == child_chunk_id,
ChildChunk.tenant_id == tenant_id
).first()
@classmethod @classmethod
def get_segments( def get_segments(
cls, cls,