|
6 | 6 | import time |
7 | 7 | import uuid |
8 | 8 | from collections import Counter |
9 | | -from typing import Any, Optional |
| 9 | +from typing import Any, Literal, Optional |
10 | 10 |
|
11 | 11 | from flask_login import current_user |
12 | 12 | from sqlalchemy import func, select |
|
51 | 51 | RetrievalModel, |
52 | 52 | SegmentUpdateArgs, |
53 | 53 | ) |
54 | | -from services.errors.account import InvalidActionError, NoPermissionError |
| 54 | +from services.errors.account import NoPermissionError |
55 | 55 | from services.errors.chunk import ChildChunkDeleteIndexError, ChildChunkIndexingError |
56 | 56 | from services.errors.dataset import DatasetNameDuplicateError |
57 | 57 | from services.errors.document import DocumentIndexingError |
@@ -1800,14 +1800,16 @@ def estimate_args_validate(cls, args: dict): |
1800 | 1800 | raise ValueError("Process rule segmentation max_tokens is invalid") |
1801 | 1801 |
|
1802 | 1802 | @staticmethod |
1803 | | - def batch_update_document_status(dataset: Dataset, document_ids: list[str], action: str, user): |
| 1803 | + def batch_update_document_status( |
| 1804 | + dataset: Dataset, document_ids: list[str], action: Literal["enable", "disable", "archive", "un_archive"], user |
| 1805 | + ): |
1804 | 1806 | """ |
1805 | 1807 | Batch update document status. |
1806 | 1808 |
|
1807 | 1809 | Args: |
1808 | 1810 | dataset (Dataset): The dataset object |
1809 | 1811 | document_ids (list[str]): List of document IDs to update |
1810 | | - action (str): Action to perform (enable, disable, archive, un_archive) |
| 1812 | + action (Literal["enable", "disable", "archive", "un_archive"]): Action to perform |
1811 | 1813 | user: Current user performing the action |
1812 | 1814 |
|
1813 | 1815 | Raises: |
@@ -1890,9 +1892,10 @@ def batch_update_document_status(dataset: Dataset, document_ids: list[str], acti |
1890 | 1892 | raise propagation_error |
1891 | 1893 |
|
1892 | 1894 | @staticmethod |
1893 | | - def _prepare_document_status_update(document, action: str, user): |
1894 | | - """ |
1895 | | - Prepare document status update information. |
| 1895 | + def _prepare_document_status_update( |
| 1896 | + document: Document, action: Literal["enable", "disable", "archive", "un_archive"], user |
| 1897 | + ): |
| 1898 | + """Prepare document status update information. |
1896 | 1899 |
|
1897 | 1900 | Args: |
1898 | 1901 | document: Document object to update |
@@ -2355,7 +2358,9 @@ def delete_segments(cls, segment_ids: list, document: Document, dataset: Dataset |
2355 | 2358 | db.session.commit() |
2356 | 2359 |
|
2357 | 2360 | @classmethod |
2358 | | - def update_segments_status(cls, segment_ids: list, action: str, dataset: Dataset, document: Document): |
| 2361 | + def update_segments_status( |
| 2362 | + cls, segment_ids: list, action: Literal["enable", "disable"], dataset: Dataset, document: Document |
| 2363 | + ): |
2359 | 2364 | # Check if segment_ids is not empty to avoid WHERE false condition |
2360 | 2365 | if not segment_ids or len(segment_ids) == 0: |
2361 | 2366 | return |
@@ -2413,8 +2418,6 @@ def update_segments_status(cls, segment_ids: list, action: str, dataset: Dataset |
2413 | 2418 | db.session.commit() |
2414 | 2419 |
|
2415 | 2420 | disable_segments_from_index_task.delay(real_deal_segment_ids, dataset.id, document.id) |
2416 | | - else: |
2417 | | - raise InvalidActionError() |
2418 | 2421 |
|
2419 | 2422 | @classmethod |
2420 | 2423 | def create_child_chunk( |
|
0 commit comments