chore: tidy up

Author: iziang

aperag/db/models.py

@@ -86,12 +86,12 @@ class DocumentIndexType(str, Enum):
 class DocumentIndexStatus(str, Enum):
     """Document index lifecycle status"""
 
-    PENDING = "PENDING"                     # Awaiting processing (create/update)
-    CREATING = "CREATING"                   # Task claimed, creation/update in progress
-    ACTIVE = "ACTIVE"                       # Index is up-to-date and ready for use
-    DELETING = "DELETING"                   # Deletion has been requested
+    PENDING = "PENDING"  # Awaiting processing (create/update)
+    CREATING = "CREATING"  # Task claimed, creation/update in progress
+    ACTIVE = "ACTIVE"  # Index is up-to-date and ready for use
+    DELETING = "DELETING"  # Deletion has been requested
     DELETION_IN_PROGRESS = "DELETION_IN_PROGRESS"  # Task claimed, deletion in progress
-    FAILED = "FAILED"                       # The last operation failed
+    FAILED = "FAILED"  # The last operation failed
 
 
 class BotStatus(str, Enum):
@@ -240,7 +240,9 @@ def get_overall_index_status(self, session) -> "DocumentStatus":
 
         if any(status == DocumentIndexStatus.FAILED for status in statuses):
             return DocumentStatus.FAILED
-        elif any(status in [DocumentIndexStatus.CREATING, DocumentIndexStatus.DELETION_IN_PROGRESS] for status in statuses):
+        elif any(
+            status in [DocumentIndexStatus.CREATING, DocumentIndexStatus.DELETION_IN_PROGRESS] for status in statuses
+        ):
             return DocumentStatus.RUNNING
         elif all(status == DocumentIndexStatus.ACTIVE for status in statuses):
             return DocumentStatus.COMPLETE

aperag/index/manager.py

@@ -12,14 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import json
 import logging
 from typing import List, Optional
 
 from sqlalchemy import and_, select
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from aperag.db.models import DocumentIndex, DocumentIndexType, DocumentIndexStatus, utc_now
+from aperag.db.models import DocumentIndex, DocumentIndexStatus, DocumentIndexType, utc_now
 
 logger = logging.getLogger(__name__)
 

aperag/index/reconciler.py

@@ -13,23 +13,22 @@
 # limitations under the License.
 
 import logging
-import time
 from typing import List, Optional
 
-from sqlalchemy import and_, or_, select, update
+from sqlalchemy import and_, select, update
 from sqlalchemy.orm import Session
 
 from aperag.config import get_sync_session
 from aperag.db.models import (
     Document,
     DocumentIndex,
-    DocumentIndexType,
     DocumentIndexStatus,
+    DocumentIndexType,
     DocumentStatus,
 )
 from aperag.tasks.scheduler import TaskScheduler, create_task_scheduler
-from aperag.utils.utils import utc_now
 from aperag.utils.constant import IndexAction
+from aperag.utils.utils import utc_now
 
 logger = logging.getLogger(__name__)
 
@@ -89,7 +88,7 @@ def _get_indexes_needing_reconciliation(self, session: Session) -> List[Document
                 DocumentIndex.status == DocumentIndexStatus.DELETING,
             ),
         }
-        
+
         for action, condition in conditions.items():
             stmt = select(DocumentIndex).where(condition)
             result = session.execute(stmt)
@@ -128,7 +127,7 @@ def _claim_document_indexes(self, session: Session, document_id: str, indexes_to
         Returns list of successfully claimed indexes with their details.
         """
         claimed_indexes = []
-        
+
         try:
             for index_id, index_type, action in indexes_to_claim:
                 if action in [IndexAction.CREATE, IndexAction.UPDATE]:
@@ -142,7 +141,7 @@ def _claim_document_indexes(self, session: Session, document_id: str, indexes_to
                 stmt = select(DocumentIndex).where(DocumentIndex.id == index_id)
                 result = session.execute(stmt)
                 current_index = result.scalar_one_or_none()
-                
+
                 if not current_index:
                     continue
 
@@ -177,13 +176,17 @@ def _claim_document_indexes(self, session: Session, document_id: str, indexes_to
                 result = session.execute(update_stmt)
                 if result.rowcount > 0:
                     # Successfully claimed this index
-                    claimed_indexes.append({
-                        'index_id': index_id,
-                        'document_id': document_id,
-                        'index_type': index_type,
-                        'action': action,
-                        'target_version': current_index.version if action in [IndexAction.CREATE, IndexAction.UPDATE] else None,
-                    })
+                    claimed_indexes.append(
+                        {
+                            "index_id": index_id,
+                            "document_id": document_id,
+                            "index_type": index_type,
+                            "action": action,
+                            "target_version": current_index.version
+                            if action in [IndexAction.CREATE, IndexAction.UPDATE]
+                            else None,
+                        }
+                    )
                     logger.debug(f"Claimed index {index_id} for document {document_id} ({action})")
                 else:
                     logger.debug(f"Could not claim index {index_id} for document {document_id}")
@@ -199,64 +202,57 @@ def _reconcile_document_operations(self, document_id: str, claimed_indexes: List
         Reconcile operations for a single document, batching same operation types together
         """
         from collections import defaultdict
-        
+
         # Group by operation type to batch operations
         operations_by_type = defaultdict(list)
         for claimed_index in claimed_indexes:
-            action = claimed_index['action']
+            action = claimed_index["action"]
             operations_by_type[action].append(claimed_index)
 
         # Process create operations as a batch
         if IndexAction.CREATE in operations_by_type:
             create_indexes = operations_by_type[IndexAction.CREATE]
-            create_types = [claimed_index['index_type'] for claimed_index in create_indexes]
+            create_types = [claimed_index["index_type"] for claimed_index in create_indexes]
             context = {}
-            
+
             for claimed_index in create_indexes:
-                index_type = claimed_index['index_type']
-                target_version = claimed_index.get('target_version')
-                
+                index_type = claimed_index["index_type"]
+                target_version = claimed_index.get("target_version")
+
                 # Store version info in context
                 if target_version is not None:
                     context[f"{index_type}_version"] = target_version
-            
-            task_id = self.task_scheduler.schedule_create_index(
-                document_id=document_id,
-                index_types=create_types,
-                context=context
+
+            self.task_scheduler.schedule_create_index(
+                document_id=document_id, index_types=create_types, context=context
             )
             logger.info(f"Scheduled create task for document {document_id}, types: {create_types}")
-        
+
         # Process update operations as a batch
         if IndexAction.UPDATE in operations_by_type:
             update_indexes = operations_by_type[IndexAction.UPDATE]
-            update_types = [claimed_index['index_type'] for claimed_index in update_indexes]
+            update_types = [claimed_index["index_type"] for claimed_index in update_indexes]
             context = {}
 
             for claimed_index in update_indexes:
-                index_type = claimed_index['index_type']
-                target_version = claimed_index.get('target_version')
-                
+                index_type = claimed_index["index_type"]
+                target_version = claimed_index.get("target_version")
+
                 # Store version info in context
                 if target_version is not None:
                     context[f"{index_type}_version"] = target_version
-            
-            task_id = self.task_scheduler.schedule_update_index(
-                document_id=document_id,
-                index_types=update_types,
-                context=context
+
+            self.task_scheduler.schedule_update_index(
+                document_id=document_id, index_types=update_types, context=context
             )
             logger.info(f"Scheduled update task for document {document_id}, types: {update_types}")
 
         # Process delete operations as a batch
         if IndexAction.DELETE in operations_by_type:
             delete_indexes = operations_by_type[IndexAction.DELETE]
-            delete_types = [claimed_index['index_type'] for claimed_index in delete_indexes]
-            
-            task_id = self.task_scheduler.schedule_delete_index(
-                document_id=document_id,
-                index_types=delete_types
-            )
+            delete_types = [claimed_index["index_type"] for claimed_index in delete_indexes]
+
+            self.task_scheduler.schedule_delete_index(document_id=document_id, index_types=delete_types)
             logger.info(f"Scheduled delete task for document {document_id}, types: {delete_types}")
 
 
@@ -322,7 +318,9 @@ def on_index_failed(document_id: str, index_type: str, error_message: str):
                         DocumentIndex.document_id == document_id,
                         DocumentIndex.index_type == DocumentIndexType(index_type),
                         # Allow transition from any in-progress state
-                        DocumentIndex.status.in_([DocumentIndexStatus.CREATING, DocumentIndexStatus.DELETION_IN_PROGRESS]),
+                        DocumentIndex.status.in_(
+                            [DocumentIndexStatus.CREATING, DocumentIndexStatus.DELETION_IN_PROGRESS]
+                        ),
                     )
                 )
                 .values(
@@ -350,15 +348,12 @@ def on_index_deleted(document_id: str, index_type: str):
         for session in get_sync_session():
             # Delete the record entirely
             from sqlalchemy import delete
-            
-            delete_stmt = (
-                delete(DocumentIndex)
-                .where(
-                    and_(
-                        DocumentIndex.document_id == document_id,
-                        DocumentIndex.index_type == DocumentIndexType(index_type),
-                        DocumentIndex.status == DocumentIndexStatus.DELETION_IN_PROGRESS,
-                    )
+
+            delete_stmt = delete(DocumentIndex).where(
+                and_(
+                    DocumentIndex.document_id == document_id,
+                    DocumentIndex.index_type == DocumentIndexType(index_type),
+                    DocumentIndex.status == DocumentIndexStatus.DELETION_IN_PROGRESS,
                 )
             )
 

aperag/service/document_service.py

@@ -15,7 +15,7 @@
 import json
 import logging
 import os
-from typing import List, Optional
+from typing import List
 
 from asgiref.sync import sync_to_async
 from fastapi import UploadFile
@@ -38,7 +38,7 @@
 from aperag.index.manager import document_index_manager
 from aperag.objectstore.base import get_object_store
 from aperag.schema import view_models
-from aperag.schema.view_models import Document, DocumentList
+from aperag.schema.view_models import DocumentList
 from aperag.utils.constant import QuotaType
 from aperag.utils.uncompress import SUPPORTED_COMPRESSED_EXTENSIONS
 
@@ -83,7 +83,8 @@ async def build_document_response(
         """Build Document response object for API return using new status model."""
         from sqlalchemy import select
 
-        from aperag.db.models import DocumentIndex, DocumentIndexType
+        from aperag.db.models import DocumentIndex
+
         # Get all document indexes for status calculation
         document_indexes = await session.execute(
             select(DocumentIndex).where(
@@ -97,12 +98,16 @@ async def build_document_response(
         # Map index states to API response format
         index_status = {}
         index_updated = {}
-        
+
         # Initialize all types as SKIPPED (when no record exists)
-        all_types = [db_models.DocumentIndexType.VECTOR, db_models.DocumentIndexType.FULLTEXT, db_models.DocumentIndexType.GRAPH]
+        all_types = [
+            db_models.DocumentIndexType.VECTOR,
+            db_models.DocumentIndexType.FULLTEXT,
+            db_models.DocumentIndexType.GRAPH,
+        ]
         for index_type in all_types:
             index_status[index_type] = "SKIPPED"
-        
+
         # Update with actual states from database
         for index in indexes:
             index_status[index.index_type] = index.status
@@ -208,9 +213,7 @@ async def _create_documents_atomically(session):
 
                     # Use index manager to create indexes with new status model
                     await document_index_manager.create_or_update_document_indexes(
-                        document_id=document_instance.id,
-                        index_types=index_types,
-                        session=session
+                        document_id=document_instance.id, index_types=index_types, session=session
                     )
 
                     # Build response object
@@ -263,14 +266,12 @@ async def _delete_document(self, session: AsyncSession, user: str, collection_id
             return
 
         # Use index manager to mark all related indexes for deletion
-        await document_index_manager.delete_document_indexes(
-            document_id=document.id, index_types=None, session=session
-        )
+        await document_index_manager.delete_document_indexes(document_id=document.id, index_types=None, session=session)
 
         # Delete from object store
         obj_store = get_object_store()
         metadata = json.loads(document.doc_metadata) if document.doc_metadata else {}
-        if object_path := metadata.get("object_path"):
+        if metadata.get("object_path"):
             try:
                 # Use delete_objects_by_prefix to remove all related files (original, chunks, etc.)
                 await sync_to_async(obj_store.delete_objects_by_prefix)(document.object_store_base_path())

aperag/tasks/scheduler.py

@@ -16,8 +16,6 @@
 from abc import ABC, abstractmethod
 from typing import Any, List, Optional
 
-from aperag.tasks.utils import cleanup_local_document, parse_document_content
-
 logger = logging.getLogger(__name__)