src/deriver/consumer.py

import logging
import time

import sentry_sdk
from pydantic import ValidationError
from sqlalchemy import select

from src import crud, models
from src.dependencies import tracked_db
from src.deriver.deriver import process_representation_tasks_batch
from src.dreamer import process_dream
from src.exceptions import ResourceNotFoundException
from src.models import Message
from src.reconciler.queue_cleanup import cleanup_queue_items
from src.reconciler.sync_vectors import run_vector_reconciliation_cycle
from src.schemas import ReconcilerType, ResolvedConfiguration
from src.telemetry.events import (
    CleanupStaleItemsCompletedEvent,
    DeletionCompletedEvent,
    SyncVectorsCompletedEvent,
    emit,
)
from src.telemetry.logging import log_performance_metrics
from src.utils import summarizer
from src.utils.queue_payload import (
    DeletionPayload,
    DreamPayload,
    ReconcilerPayload,
    SummaryPayload,
    WebhookPayload,
)
from src.webhooks import webhook_delivery

logger = logging.getLogger(__name__)
logging.getLogger("sqlalchemy.engine.Engine").disabled = True


async def process_item(queue_item: models.QueueItem) -> None:
    """Process a single item from the queue."""
    task_type = queue_item.task_type
    queue_payload = queue_item.payload
    workspace_name = queue_item.workspace_name

    # Handle reconciler first - it's the only task type that doesn't require workspace_name
    if task_type == "reconciler":
        with sentry_sdk.start_transaction(name="process_reconciler_task", op="deriver"):
            try:
                validated = ReconcilerPayload(**queue_payload)
            except ValidationError as e:
                logger.error(
                    "Invalid reconciler payload received: %s. Payload: %s",
                    str(e),
                    queue_payload,
                )
                raise ValueError(f"Invalid payload structure: {str(e)}") from e
            await process_reconciler(validated)
        return

    # All other task types require a workspace_name
    if workspace_name is None:
        raise ValueError(f"{task_type} tasks require a workspace_name")

    if task_type == "webhook":
        try:
            validated = WebhookPayload(**queue_payload)
        except ValidationError as e:
            logger.error(
                "Invalid webhook payload received: %s. Payload: %s",
                str(e),
                queue_payload,
            )
            raise ValueError(f"Invalid payload structure: {str(e)}") from e
        await webhook_delivery.deliver_webhook(validated, workspace_name)

    elif task_type == "summary":
        try:
            validated = SummaryPayload(**queue_payload)
        except ValidationError as e:
            logger.error(
                "Invalid summary payload received: %s. Payload: %s",
                str(e),
                queue_payload,
            )
            raise ValueError(f"Invalid payload structure: {str(e)}") from e

        # Use workspace_name and message_id from QueueItem columns
        message_id = queue_item.message_id

        if message_id is None:
            raise ValueError("Summary tasks require a message_id")

        message_public_id = validated.message_public_id
        if not message_public_id:
            logger.debug("Fetching message public ID for message %s", message_id)
            async with tracked_db(operation_name="summary_fallback") as db:
                stmt = (
                    select(models.Message)
                    .where(models.Message.workspace_name == workspace_name)
                    .where(models.Message.session_name == validated.session_name)
                    .where(models.Message.id == message_id)
                )
                result = await db.execute(stmt)

                message = result.scalar_one_or_none()
                if message is None:
                    logger.error(
                        "Failed to fetch message with ID %s for process_summary_task",
                        message_id,
                    )
                    return
                message_public_id = message.public_id

        with sentry_sdk.start_transaction(name="process_summary_task", op="deriver"):
            await summarizer.summarize_if_needed(
                workspace_name,
                validated.session_name,
                message_id,
                validated.message_seq_in_session,
                message_public_id,
                validated.configuration,
            )
            log_performance_metrics("summary", f"{workspace_name}_{message_id}")

    elif task_type == "dream":
        with sentry_sdk.start_transaction(name="process_dream_task", op="deriver"):
            try:
                validated = DreamPayload(**queue_payload)
            except ValidationError as e:
                logger.error(
                    "Invalid dream payload received: %s. Payload: %s",
                    str(e),
                    queue_payload,
                )
                raise ValueError(f"Invalid payload structure: {str(e)}") from e
            await process_dream(validated, workspace_name)

    elif task_type == "deletion":
        with sentry_sdk.start_transaction(name="process_deletion_task", op="deriver"):
            try:
                validated = DeletionPayload(**queue_payload)
            except ValidationError as e:
                logger.error(
                    "Invalid deletion payload received: %s. Payload: %s",
                    str(e),
                    queue_payload,
                )
                raise ValueError(f"Invalid payload structure: {str(e)}") from e
            await process_deletion(validated, workspace_name)

    else:
        raise ValueError(f"Invalid task type: {task_type}")


async def process_representation_batch(
    messages: list[Message],
    message_level_configuration: ResolvedConfiguration | None,
    *,
    observers: list[str] | None,
    observed: str | None,
    queue_item_message_ids: list[int],
) -> None:
    """
    Prepares and processes a batch of messages for representation tasks.

    Args:
        messages: List of messages to process
        message_level_configuration: Resolved configuration for this batch
        observers: List of observers for the messages
        observed: The observed of the messages
        queue_item_message_ids: Message IDs from queue items
    """
    if not messages or not messages[0]:
        logger.debug("process_representation_batch received no messages")
        return

    if observed is None or observers is None or len(observers) == 0:
        raise ValueError("observed and observers are required for representation tasks")

    await process_representation_tasks_batch(
        messages,
        message_level_configuration,
        observers=observers,
        observed=observed,
        queue_item_message_ids=queue_item_message_ids,
    )


async def process_deletion(
    payload: DeletionPayload,
    workspace_name: str,
) -> None:
    """
    Process a deletion task from the queue.

    This function handles the actual deletion of resources based on the deletion type.
    It is designed to be idempotent - deleting an already-deleted resource is a no-op.

    Args:
        payload: The deletion payload containing deletion_type and resource_id
        workspace_name: The workspace name for scoping the deletion

    Raises:
        ValueError: If the deletion type is not supported
    """
    deletion_type = payload.deletion_type
    resource_id = payload.resource_id
    success = True
    error_message: str | None = None
    peers_deleted = 0
    sessions_deleted = 0
    messages_deleted = 0
    conclusions_deleted = 0

    logger.info(
        "Processing deletion task: type=%s, resource_id=%s, workspace=%s",
        deletion_type,
        resource_id,
        workspace_name,
    )

    async with tracked_db("process_deletion") as db:
        if deletion_type == "session":
            try:
                result = await crud.delete_session(
                    db, workspace_name=workspace_name, session_name=resource_id
                )
                messages_deleted = result.messages_deleted
                conclusions_deleted = result.conclusions_deleted
                logger.info(
                    "Successfully deleted session %s in workspace %s "
                    + "(messages=%d, conclusions=%d)",
                    resource_id,
                    workspace_name,
                    messages_deleted,
                    conclusions_deleted,
                )
            except ResourceNotFoundException as e:
                # Session not found - may have already been deleted, treat as success
                logger.warning(
                    "Session %s not found during deletion (may already be deleted): %s",
                    resource_id,
                    str(e),
                )

        elif deletion_type == "observation":
            try:
                await crud.delete_document_by_id(
                    db, workspace_name=workspace_name, document_id=resource_id
                )
                conclusions_deleted = 1  # Single observation deleted
                logger.info(
                    "Successfully deleted observation %s in workspace %s",
                    resource_id,
                    workspace_name,
                )
            except ResourceNotFoundException as e:
                # Document not found - may have already been deleted, treat as success
                logger.warning(
                    "Observation %s not found during deletion (may already be deleted): %s",
                    resource_id,
                    str(e),
                )

        elif deletion_type == "workspace":
            try:
                result = await crud.delete_workspace(db, workspace_name=workspace_name)
                peers_deleted = result.peers_deleted
                sessions_deleted = result.sessions_deleted
                messages_deleted = result.messages_deleted
                conclusions_deleted = result.conclusions_deleted
                logger.info(
                    "Successfully deleted workspace %s "
                    + "(peers=%d, sessions=%d, messages=%d, conclusions=%d)",
                    workspace_name,
                    peers_deleted,
                    sessions_deleted,
                    messages_deleted,
                    conclusions_deleted,
                )
            except ResourceNotFoundException as e:
                # Workspace not found - may have already been deleted, treat as success
                logger.warning(
                    "Workspace %s not found during deletion (may already be deleted): %s",
                    workspace_name,
                    str(e),
                )

        else:
            success = False
            error_message = f"Unsupported deletion type: {deletion_type}"
            raise ValueError(error_message)

    # Emit telemetry event
    emit(
        DeletionCompletedEvent(
            workspace_name=workspace_name,
            deletion_type=deletion_type,
            resource_id=resource_id,
            success=success,
            peers_deleted=peers_deleted,
            sessions_deleted=sessions_deleted,
            messages_deleted=messages_deleted,
            conclusions_deleted=conclusions_deleted,
            error_message=error_message,
        )
    )


async def process_reconciler(payload: ReconcilerPayload) -> None:
    """
    Process a reconciler task from the queue.

    Currently supports:
    - sync_vectors: Syncs pending documents/message embeddings to vector store
      and cleans up soft-deleted documents.
    - cleanup_queue: Removes old processed queue items.

    Args:
        payload: The reconciler payload containing the reconciler type
    """
    reconciler_type = payload.reconciler_type
    start_time = time.perf_counter()

    if reconciler_type == ReconcilerType.SYNC_VECTORS:
        logger.debug("Processing sync_vectors task")
        metrics = await run_vector_reconciliation_cycle()

        duration_ms = (time.perf_counter() - start_time) * 1000

        if (
            metrics.total_synced > 0
            or metrics.total_failed > 0
            or metrics.total_cleaned > 0
        ):
            logger.info(
                "Reconciliation complete: synced %s docs, %s message embeddings; failed %s docs, %s message embeddings; cleaned %s docs",
                metrics.documents_synced,
                metrics.message_embeddings_synced,
                metrics.documents_failed,
                metrics.message_embeddings_failed,
                metrics.documents_cleaned,
            )

            # Emit telemetry event
            emit(
                SyncVectorsCompletedEvent(
                    documents_synced=metrics.documents_synced,
                    documents_failed=metrics.documents_failed,
                    documents_cleaned=metrics.documents_cleaned,
                    message_embeddings_synced=metrics.message_embeddings_synced,
                    message_embeddings_failed=metrics.message_embeddings_failed,
                    total_duration_ms=duration_ms,
                )
            )

    elif reconciler_type == ReconcilerType.CLEANUP_QUEUE:
        logger.debug("Processing cleanup_queue task")
        deleted_count = await cleanup_queue_items()

        duration_ms = (time.perf_counter() - start_time) * 1000

        if deleted_count > 0:
            # Emit telemetry event for cleanup stale items
            emit(
                CleanupStaleItemsCompletedEvent(
                    total_duration_ms=duration_ms,
                )
            )
    else:
        raise ValueError(f"Unsupported reconciler type: {reconciler_type}")