396 lines
15 KiB
Python
396 lines
15 KiB
Python
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, ValidationException
|
|
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],
|
|
hit_batch_token_cap: bool = False,
|
|
was_flush_enabled: bool = False,
|
|
batch_max_tokens: int = 0,
|
|
) -> 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
|
|
hit_batch_token_cap: whether the queue batcher clamped this batch to fit
|
|
was_flush_enabled: snapshot of DERIVER.FLUSH_ENABLED at fetch time
|
|
batch_max_tokens: DERIVER.REPRESENTATION_BATCH_MAX_TOKENS snapshot
|
|
"""
|
|
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,
|
|
hit_batch_token_cap=hit_batch_token_cap,
|
|
was_flush_enabled=was_flush_enabled,
|
|
batch_max_tokens=batch_max_tokens,
|
|
)
|
|
|
|
|
|
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,
|
|
)
|
|
|
|
# try/except/finally so the event ALWAYS fires — both success and
|
|
# failure (unsupported type, unexpected CRUD error). Previously the
|
|
# unsupported-type branch raised before the emit ran, and unexpected
|
|
# CRUD errors bubbled up without telemetry.
|
|
try:
|
|
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 ValidationException(error_message)
|
|
except Exception as e:
|
|
# Catch anything that survived the per-branch `ResourceNotFoundException`
|
|
# handling above (incl. the ValueError from the unsupported-type branch).
|
|
# Record telemetry, then re-raise so the queue worker still surfaces
|
|
# the failure to its caller.
|
|
success = False
|
|
if error_message is None:
|
|
error_message = f"{type(e).__name__}: {e}"
|
|
raise
|
|
finally:
|
|
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(
|
|
queue_items_cleaned=deleted_count,
|
|
total_duration_ms=duration_ms,
|
|
)
|
|
)
|
|
else:
|
|
raise ValueError(f"Unsupported reconciler type: {reconciler_type}")
|