diff --git a/api/controllers/console/datasets/rag_pipeline/rag_pipeline_datasets.py b/api/controllers/console/datasets/rag_pipeline/rag_pipeline_datasets.py index 3bbbecaecfd7f6..094adc42022799 100644 --- a/api/controllers/console/datasets/rag_pipeline/rag_pipeline_datasets.py +++ b/api/controllers/console/datasets/rag_pipeline/rag_pipeline_datasets.py @@ -17,7 +17,7 @@ with_current_user, ) from extensions.ext_database import db -from fields.dataset_fields import DatasetDetailResponse +from fields.dataset_fields import DatasetDetailResponse, dataset_detail_response_source from libs.helper import dump_response from libs.login import login_required from models import Account @@ -116,6 +116,7 @@ def post(self, current_tenant_id: str, current_user: Account) -> JsonResponseWit # The role of the current user in the ta table must be admin, owner, or editor, or dataset_operator if not current_user.is_dataset_editor: raise Forbidden() + session = db.session() dataset = DatasetService.create_empty_rag_pipeline_dataset( tenant_id=current_tenant_id, rag_pipeline_dataset_create_entity=RagPipelineDatasetCreateEntity( @@ -129,6 +130,6 @@ def post(self, current_tenant_id: str, current_user: Account) -> JsonResponseWit permission=DatasetPermissionEnum.ONLY_ME, partial_member_list=None, ), - session=db.session(), + session=session, ) - return dump_response(DatasetDetailResponse, dataset), 201 + return dump_response(DatasetDetailResponse, dataset_detail_response_source(dataset, session=session)), 201 diff --git a/api/models/dataset.py b/api/models/dataset.py index 891d949c8e6901..6fe941bb61630f 100644 --- a/api/models/dataset.py +++ b/api/models/dataset.py @@ -212,17 +212,9 @@ class Dataset(Base): enable_api = mapped_column(sa.Boolean, nullable=False, server_default=sa.text("true")) is_multimodal = mapped_column(sa.Boolean, default=False, nullable=False, server_default=sa.text("false")) - @property - def total_documents(self) -> int: - return self.get_total_documents(session=db.session()) - def get_total_documents(self, *, session: Session) -> int: return self.get_document_count(session=session) - @property - def total_available_documents(self) -> int: - return self.get_total_available_documents(session=db.session()) - def get_total_available_documents(self, *, session: Session) -> int: return ( session.scalar( @@ -258,20 +250,12 @@ def created_by_account(self): def get_created_by_account(self, *, session: Session) -> Account | None: return session.get(Account, self.created_by) - @property - def author_name(self) -> str | None: - return self.get_author_name(session=db.session()) - def get_author_name(self, *, session: Session) -> str | None: account = self.get_created_by_account(session=session) if account: return account.name return None - @property - def latest_process_rule(self): - return self.get_latest_process_rule(session=db.session()) - def get_latest_process_rule(self, *, session: Session) -> "DatasetProcessRule | None": return session.scalar( select(DatasetProcessRule) @@ -391,10 +375,6 @@ def get_tags(self, *, session: Session) -> Sequence[Tag]: return tags or [] - @property - def external_knowledge_info(self) -> dict[str, Any] | None: - return self.get_external_knowledge_info(session=db.session()) - def get_external_knowledge_info(self, *, session: Session) -> dict[str, Any] | None: if self.provider != "external": return None diff --git a/api/tests/test_containers_integration_tests/models/test_dataset_models.py b/api/tests/test_containers_integration_tests/models/test_dataset_models.py index a3bbf196576303..ecfc51b12a2ffc 100644 --- a/api/tests/test_containers_integration_tests/models/test_dataset_models.py +++ b/api/tests/test_containers_integration_tests/models/test_dataset_models.py @@ -49,7 +49,7 @@ def test_dataset_with_documents_relationship(self, db_session_with_containers: S db_session_with_containers.add(doc) db_session_with_containers.flush() - assert dataset.total_documents == 3 + assert dataset.get_total_documents(session=db_session_with_containers) == 3 def test_dataset_available_documents_count(self, db_session_with_containers: Session) -> None: """Test dataset can count available documents.""" @@ -104,7 +104,7 @@ def test_dataset_available_documents_count(self, db_session_with_containers: Ses db_session_with_containers.add_all([doc_available, doc_pending, doc_disabled]) db_session_with_containers.flush() - assert dataset.total_available_documents == 1 + assert dataset.get_total_available_documents(session=db_session_with_containers) == 1 def test_dataset_word_count_aggregation(self, db_session_with_containers: Session) -> None: """Test dataset can aggregate word count from documents."""