copilot-swe-agent[bot]
8d7c8e7c4e
- Add embedding_model config setting (replaces hardcoded text-embedding-3-small) - Add compute_document_embedding Celery task for ingestion-time embedding - Chain embedding task into finalize_document_storage pipeline - Add backfill_missing_embeddings periodic task (every 5 min) for legacy files - Add debug API endpoints: embedding-status, compute-embedding, diagnostic/embeddings, diagnostic/compute-all-embeddings - Refactor find_similar_documents to only use pre-computed embeddings (no lazy API calls) - Use yield_per(500) and column-only queries for 100K+ scale - Add embedding status indicator and recompute button in file detail UI Co-authored-by: christianlouis <361235+christianlouis@users.noreply.github.com>
104 lines
4.3 KiB
Python
104 lines
4.3 KiB
Python
#!/usr/bin/env python3
|
|
|
|
import logging
|
|
import os
|
|
|
|
# Import the shared Celery instance
|
|
from app.celery_app import celery
|
|
from app.config import settings
|
|
from app.database import SessionLocal
|
|
from app.models import FileRecord
|
|
from app.tasks.retry_config import BaseTaskWithRetry
|
|
|
|
# Import the aggregator task and validator
|
|
from app.tasks.send_to_all import get_configured_services_from_validator, send_to_all_destinations
|
|
|
|
# Import database and logging utils from main
|
|
from app.utils import log_task_progress
|
|
|
|
# Import notification utility
|
|
from app.utils.notification import notify_file_processed
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@celery.task(base=BaseTaskWithRetry, bind=True)
|
|
def finalize_document_storage(self, original_file: str, processed_file: str, metadata: dict, file_id: int = None):
|
|
"""
|
|
Final storage step after embedding metadata.
|
|
We will now call 'send_to_all_destinations' to push the final PDF to Dropbox/Nextcloud/Paperless.
|
|
After uploading, send a notification about the processed file.
|
|
"""
|
|
task_id = self.request.id
|
|
logger.info(f"[{task_id}] Finalizing document storage for {processed_file}")
|
|
|
|
# 1. Update Database Status (From Main)
|
|
log_task_progress(
|
|
task_id,
|
|
"finalize_document_storage",
|
|
"in_progress",
|
|
f"Finalizing: {os.path.basename(processed_file)}",
|
|
file_id=file_id,
|
|
)
|
|
|
|
# Get file_id from database if not provided (fallback logic from Main)
|
|
if file_id is None:
|
|
with SessionLocal() as db:
|
|
# Only as a last resort, try to find by exact match on local_filename
|
|
tmp_path = os.path.join(settings.workdir, "tmp", os.path.basename(original_file))
|
|
file_record = db.query(FileRecord).filter(FileRecord.local_filename == tmp_path).first()
|
|
if file_record:
|
|
file_id = file_record.id
|
|
|
|
# 2. Determine Configured Destinations (From Copilot)
|
|
# This is needed for the notification message later
|
|
configured_destinations = []
|
|
try:
|
|
configured_services = get_configured_services_from_validator()
|
|
# Get list of service names that are configured
|
|
for service_name, is_configured in configured_services.items():
|
|
if is_configured:
|
|
# Format service names for display
|
|
display_name = service_name.replace("_", " ").title()
|
|
configured_destinations.append(display_name)
|
|
except Exception as e:
|
|
logger.warning(f"[WARNING] Could not determine configured destinations: {e}")
|
|
configured_destinations = ["configured destinations"]
|
|
|
|
# 3. Queue Uploads (Merged)
|
|
# Uses Main branch signature to ensure file_id is passed, but keeps logic structure
|
|
logger.info(f"[{task_id}] Queueing uploads to all destinations")
|
|
log_task_progress(
|
|
task_id, "finalize_document_storage", "success", "Queuing uploads to destinations", file_id=file_id
|
|
)
|
|
|
|
# Note: send_to_all_destinations is asynchronous and queues upload tasks
|
|
# We pass 'True' (delete_after) and 'file_id' as per Main branch requirements
|
|
send_to_all_destinations.delay(processed_file, True, file_id)
|
|
|
|
# 3a. Queue embedding computation so similarity scores are ready for queries
|
|
if file_id is not None:
|
|
try:
|
|
from app.tasks.compute_embedding import compute_document_embedding
|
|
|
|
compute_document_embedding.delay(file_id)
|
|
logger.info(f"[{task_id}] Queued embedding computation for file {file_id}")
|
|
except Exception as e:
|
|
logger.warning(f"[{task_id}] Could not queue embedding task: {e}")
|
|
|
|
# 4. Send Notification (From Copilot)
|
|
# Note: This notification is sent after processing is complete but while uploads
|
|
# are being queued.
|
|
try:
|
|
# Get file information
|
|
file_size = os.path.getsize(processed_file) if os.path.exists(processed_file) else 0
|
|
filename = os.path.basename(processed_file)
|
|
|
|
notify_file_processed(
|
|
filename=filename, file_size=file_size, metadata=metadata, destinations=configured_destinations
|
|
)
|
|
except Exception as e:
|
|
logger.warning(f"[WARNING] Failed to send file processed notification: {e}")
|
|
|
|
return {"status": "Completed", "file": processed_file}
|