replypier/backend/app/repositories/knowledge_queue.py

"""
Copyright 2024-2026 ReplyPier

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

from sqlalchemy import func
from sqlalchemy.orm import Session
from app.models.knowledge_queue import KnowledgeQueue, QueueStatus
from typing import List, Optional
from datetime import datetime
from app.core.logger import get_logger
from app.core.logger import get_logger

logger = get_logger(__name__)


class KnowledgeQueueRepository:
    def __init__(self, db: Session):
        self.db = db

    def create(self, queue_item: KnowledgeQueue) -> KnowledgeQueue:
        self.db.add(queue_item)
        self.db.commit()
        self.db.refresh(queue_item)
        return queue_item

    def get_by_id(self, queue_id: int) -> Optional[KnowledgeQueue]:
        """Get queue item by ID"""
        return self.db.query(KnowledgeQueue).filter(
            KnowledgeQueue.id == queue_id
        ).first()

    def get_pending(self) -> List[KnowledgeQueue]:
        """Get all pending queue items ordered by priority (highest first), then by creation time"""
        return self.db.query(KnowledgeQueue)\
            .filter(KnowledgeQueue.status == QueueStatus.PENDING)\
            .order_by(KnowledgeQueue.priority.desc(), KnowledgeQueue.created_at.asc())\
            .all()

    def get_abandoned(self, last_progress_before: datetime) -> List[KnowledgeQueue]:
        """Items still 'processing' with no progress since `last_progress_before`,
        locked so two workers never handle the same one."""
        return self.db.query(KnowledgeQueue)\
            .filter(KnowledgeQueue.status == QueueStatus.PROCESSING,
                    func.coalesce(KnowledgeQueue.updated_at, KnowledgeQueue.created_at) < last_progress_before)\
            .with_for_update(skip_locked=True)\
            .all()

    def update_status(self, queue_id: int, status: QueueStatus, error: Optional[str] = None) -> bool:
        item = self.db.query(KnowledgeQueue).filter(
            KnowledgeQueue.id == queue_id).first()
        if item:
            item.status = status
            item.error = error
            # Update processing stage to COMPLETED when status is COMPLETED
            if status == QueueStatus.COMPLETED:
                item.processing_stage = "COMPLETED"
                item.progress_percentage = 100.0
            self.db.commit()
            return True
        return False
        
    def update_progress(
        self, 
        queue_id: int, 
        processing_stage: str = None, 
        progress_percentage: float = None,
        total_items: int = None,
        processed_items: int = None,
        crawled_url: str = None,
        force_stage_update: bool = False
    ) -> bool:
        """Update the progress of a queue item"""
        item = self.db.query(KnowledgeQueue).filter(
            KnowledgeQueue.id == queue_id).first()
        if item:
            # Define stage priority to prevent regression
            stage_priority = {
                "not_started": 0,
                "crawling": 1,
                "embedding": 2,
                "completed": 3
            }
            
            # Only update processing stage if it's an advancement or forced
            if processing_stage and (force_stage_update or not item.processing_stage or 
                stage_priority.get(processing_stage.lower(), 0) >= stage_priority.get(item.processing_stage.lower(), 0)):
                item.processing_stage = processing_stage
            
            # Update progress percentage using main columns
            if progress_percentage is not None:
                item.progress_percentage = progress_percentage
            elif processed_items is not None and total_items is not None and total_items > 0:
                # Auto-calculate percentage based on processed/total items
                item.progress_percentage = (processed_items / total_items) * 100
                
            if total_items is not None:
                item.total_items = total_items
                
            if processed_items is not None:
                item.processed_items = processed_items
            
            # Add crawled URL to the list if provided (store only URL strings)
            if crawled_url:
                if not item.crawled_urls:
                    item.crawled_urls = []
                # Store only the URL string, no timestamp or status
                if crawled_url not in item.crawled_urls:
                    # Create a new list to ensure SQLAlchemy detects the change
                    new_crawled_urls = list(item.crawled_urls) if item.crawled_urls else []
                    new_crawled_urls.append(crawled_url)
                    item.crawled_urls = new_crawled_urls
                    
                    # Mark the attribute as modified to ensure SQLAlchemy commits the change
                    from sqlalchemy.orm.attributes import flag_modified
                    flag_modified(item, 'crawled_urls')
                    
                    from app.core.logger import get_logger
                    logger = get_logger(__name__)
                    logger.debug(f"📝 Added crawled URL to queue {queue_id}: {crawled_url} (total: {len(item.crawled_urls)})")
                    
                else:
                    from app.core.logger import get_logger
                    logger = get_logger(__name__)
                    logger.debug(f"🔄 URL already in list for queue {queue_id}: {crawled_url}")
                    
            # Commit changes and verify
            self.db.commit()
            
            # If we added a crawled URL, verify it was saved
            if crawled_url:
                self.db.refresh(item)
                from app.core.logger import get_logger
                logger = get_logger(__name__)
                logger.debug(f"🔍 After commit - queue {queue_id} crawled_urls count: {len(item.crawled_urls) if item.crawled_urls else 0}")
            
            return True
        return False