"""ARQ background jobs for the Unified Search plugin.""" from __future__ import annotations import logging import uuid from typing import Any from app.core.db import get_session_factory logger = logging.getLogger(__name__) BATCH_SIZE = 100 # Entity type -> table name mapping (shared by multiple jobs) _TABLE_MAP: dict[str, str] = { "contact": "contacts", "mail": "mails", "file": "files", "event": "calendar_entries", } def _parse_id(id_str: str) -> uuid.UUID: """Parse a string to UUID.""" if isinstance(id_str, uuid.UUID): return id_str return uuid.UUID(str(id_str)) async def index_mails(ctx: dict[str, Any], mail_ids: list[str]) -> None: """Index mails: generate and store embeddings.""" from app.plugins.builtins.unified_search.embedding import index_entity factory = get_session_factory() async with factory() as db: for mail_id in mail_ids: try: eid = _parse_id(mail_id) # tenant_id is derived from the mail itself from sqlalchemy import text result = await db.execute( text("SELECT tenant_id FROM mails WHERE id = :mid"), {"mid": eid}, ) row = result.mappings().first() if not row: continue tenant_id = row["tenant_id"] await index_entity("mail", eid, tenant_id, db) except Exception: logger.exception("Failed to index mail %s", mail_id) try: await db.rollback() except Exception: pass async def index_file(ctx: dict[str, Any], file_id: str) -> None: """Index a file: extract text, store content_text, generate embedding.""" from sqlalchemy import text from app.plugins.builtins.unified_search.embedding import generate_embedding from app.plugins.builtins.unified_search.text_extraction import extract_text_from_file factory = get_session_factory() async with factory() as db: try: eid = _parse_id(file_id) result = await db.execute( text("SELECT tenant_id, storage_path, mime_type FROM files WHERE id = :fid"), {"fid": eid}, ) row = result.mappings().first() if not row: logger.warning("File not found: %s", file_id) return tenant_id = row["tenant_id"] storage_path = row["storage_path"] mime_type = row["mime_type"] # Extract text content_text = await extract_text_from_file(storage_path, mime_type) # Store content_text await db.execute( text("UPDATE files SET content_text = :ct WHERE id = :fid"), {"ct": content_text, "fid": eid}, ) await db.commit() # Generate embedding from extracted text + filename result_name = await db.execute( text("SELECT name FROM files WHERE id = :fid"), {"fid": eid}, ) name_row = result_name.mappings().first() name = name_row["name"] if name_row else "" embedding_text = f"{name} {content_text[:5000]}" if embedding_text.strip(): embedding = await generate_embedding(embedding_text, db=db, tenant_id=tenant_id) if embedding: await db.execute( text("UPDATE files SET embedding = cast(:emb AS vector) WHERE id = :fid"), {"emb": str(embedding), "fid": eid}, ) await db.commit() logger.info("Indexed file %s", file_id) except Exception: logger.exception("Failed to index file %s", file_id) try: await db.rollback() except Exception: pass async def index_contact(ctx: dict[str, Any], contact_id: str) -> None: """Index a contact: generate and store embedding.""" from sqlalchemy import text from app.plugins.builtins.unified_search.embedding import index_entity factory = get_session_factory() async with factory() as db: try: eid = _parse_id(contact_id) result = await db.execute( text("SELECT tenant_id FROM contacts WHERE id = :cid"), {"cid": eid}, ) row = result.mappings().first() if not row: return await index_entity("contact", eid, row["tenant_id"], db) except Exception: logger.exception("Failed to index contact %s", contact_id) try: await db.rollback() except Exception: pass async def index_event(ctx: dict[str, Any], event_id: str) -> None: """Index a calendar event: generate and store embedding.""" from sqlalchemy import text from app.plugins.builtins.unified_search.embedding import index_entity factory = get_session_factory() async with factory() as db: try: eid = _parse_id(event_id) result = await db.execute( text("SELECT tenant_id FROM calendar_entries WHERE id = :eid"), {"eid": eid}, ) row = result.mappings().first() if not row: return await index_entity("event", eid, row["tenant_id"], db) except Exception: logger.exception("Failed to index event %s", event_id) try: await db.rollback() except Exception: pass async def reindex(ctx: dict[str, Any], entity_type: str) -> None: """Reindex all entities of a given type with pagination. Clears existing embeddings before re-indexing, tracks progress, and continues on individual failures. """ from sqlalchemy import text from app.plugins.builtins.unified_search.embedding import index_entity table = _TABLE_MAP.get(entity_type) if not table: logger.warning("Unknown entity_type for reindex: %s", entity_type) return factory = get_session_factory() async with factory() as db: # Clear existing embeddings before re-indexing try: await db.execute( text(f"UPDATE {table} SET embedding = NULL WHERE deleted_at IS NULL"), ) await db.commit() except Exception: logger.exception("Failed to clear embeddings for %s", entity_type) try: await db.rollback() except Exception: pass total_indexed = 0 total_failed = 0 offset = 0 while True: result = await db.execute( text( f"SELECT id, tenant_id FROM {table} " f"WHERE deleted_at IS NULL ORDER BY created_at LIMIT :lim OFFSET :off" ), {"lim": BATCH_SIZE, "off": offset}, ) rows = result.mappings().all() if not rows: break for row in rows: try: await index_entity( entity_type, row["id"], row["tenant_id"], db, ) total_indexed += 1 except Exception: logger.exception("Reindex failed for %s/%s", entity_type, row["id"]) total_failed += 1 try: await db.rollback() except Exception: pass # Progress tracking every BATCH_SIZE entities logger.info( "Reindex progress for %s: %d indexed, %d failed (offset=%d)", entity_type, total_indexed, total_failed, offset, ) offset += BATCH_SIZE logger.info( "Reindex complete for %s: %d indexed, %d failed", entity_type, total_indexed, total_failed, ) async def reindex_all(ctx: dict[str, Any]) -> None: """Reindex all entity types in sequence, including file chunks.""" entity_types = ["contact", "mail", "file", "event"] for etype in entity_types: try: await reindex(ctx, etype) except Exception: logger.exception("reindex_all: reindex failed for %s", etype) # For files, also re-index chunks if etype == "file": try: from sqlalchemy import text factory = get_session_factory() async with factory() as db: result = await db.execute( text("SELECT id FROM files WHERE deleted_at IS NULL"), ) rows = result.mappings().all() for row in rows: try: await index_file_chunks(ctx, str(row["id"])) except Exception: logger.exception("reindex_all: chunk re-index failed for file %s", row["id"]) except Exception: logger.exception("reindex_all: chunk re-index batch failed") logger.info("reindex_all complete") async def delete_entity_index(ctx: dict[str, Any], entity_type: str, entity_id: str) -> None: """Delete an entity's embedding (set embedding=NULL). Logs the action to SearchIndexLog on success. """ from sqlalchemy import text from app.plugins.builtins.unified_search.models import SearchIndexLog table = _TABLE_MAP.get(entity_type) if not table: logger.warning("delete_entity_index: unknown entity_type=%s", entity_type) return factory = get_session_factory() async with factory() as db: try: eid = _parse_id(entity_id) # Get tenant_id from the entity result = await db.execute( text(f"SELECT tenant_id FROM {table} WHERE id = :eid"), {"eid": eid}, ) row = result.mappings().first() tenant_id = row["tenant_id"] if row else None await db.execute( text( f"UPDATE {table} SET embedding = NULL, indexed_at = NULL " f"WHERE id = :eid" ), {"eid": eid}, ) await db.commit() # Log to SearchIndexLog if tenant_id: log_entry = SearchIndexLog( tenant_id=tenant_id, entity_type=entity_type, entity_id=eid, action="delete", status="success", ) db.add(log_entry) await db.commit() logger.info("Deleted index for %s/%s", entity_type, entity_id) except Exception: logger.exception("Failed to delete index for %s/%s", entity_type, entity_id) try: await db.rollback() except Exception: pass async def delete_file_chunks(ctx: dict[str, Any], file_id: str) -> None: """Delete all document_chunks for a file.""" from sqlalchemy import text factory = get_session_factory() async with factory() as db: try: eid = _parse_id(file_id) await db.execute( text("DELETE FROM document_chunks WHERE file_id = :fid"), {"fid": eid}, ) await db.commit() logger.info("Deleted chunks for file %s", file_id) except Exception: logger.exception("Failed to delete chunks for file %s", file_id) try: await db.rollback() except Exception: pass async def retry_failed_index(ctx: dict[str, Any], entity_type: str, entity_id: str) -> None: """Retry a failed index operation.""" from sqlalchemy import text from app.plugins.builtins.unified_search.embedding import index_entity table = _TABLE_MAP.get(entity_type) if not table: logger.warning("retry_failed_index: unknown entity_type=%s", entity_type) return factory = get_session_factory() async with factory() as db: try: eid = _parse_id(entity_id) result = await db.execute( text(f"SELECT tenant_id FROM {table} WHERE id = :eid"), {"eid": eid}, ) row = result.mappings().first() if not row: logger.warning("retry_failed_index: entity not found %s/%s", entity_type, entity_id) return await index_entity(entity_type, eid, row["tenant_id"], db) except Exception: logger.exception("retry_failed_index failed for %s/%s", entity_type, entity_id) try: await db.rollback() except Exception: pass async def embedding_batch(ctx: dict[str, Any]) -> None: """Periodic job: find entities without embeddings and index them.""" from sqlalchemy import text from app.plugins.builtins.unified_search.embedding import index_entity factory = get_session_factory() async with factory() as db: for etype, table in _TABLE_MAP.items(): try: result = await db.execute( text( f"SELECT id, tenant_id FROM {table} " f"WHERE deleted_at IS NULL AND embedding IS NULL " f"LIMIT :lim" ), {"lim": BATCH_SIZE}, ) rows = result.mappings().all() for row in rows: try: await index_entity(etype, row["id"], row["tenant_id"], db) except Exception: logger.exception("Batch index failed for %s/%s", etype, row["id"]) try: await db.rollback() except Exception: pass except Exception: logger.exception("Batch query failed for %s", etype) try: await db.rollback() except Exception: pass logger.info("Embedding batch job complete") async def index_file_chunks(ctx: dict[str, Any], file_id: str) -> None: """Extract text from a file, chunk it, generate embeddings, and store in document_chunks.""" from sqlalchemy import text from app.plugins.builtins.unified_search.chunking import chunk_text from app.plugins.builtins.unified_search.embedding import generate_embedding from app.plugins.builtins.unified_search.text_extraction import extract_text_from_file factory = get_session_factory() async with factory() as db: try: eid = _parse_id(file_id) result = await db.execute( text("SELECT tenant_id, storage_path, mime_type, name FROM files WHERE id = :fid"), {"fid": eid}, ) row = result.mappings().first() if not row: logger.warning("File not found for chunking: %s", file_id) return tenant_id = row["tenant_id"] storage_path = row["storage_path"] mime_type = row["mime_type"] name = row.get("name", "") or "" # Extract text from file content_text = await extract_text_from_file(storage_path, mime_type) if not content_text.strip(): logger.debug("No text extracted from file %s, skipping chunking", file_id) return # Store content_text on the file record await db.execute( text("UPDATE files SET content_text = :ct WHERE id = :fid"), {"ct": content_text, "fid": eid}, ) await db.commit() # Chunk the text (include filename for context) full_text = f"{name}\n{content_text}" chunks = chunk_text(full_text, chunk_size=1000, overlap=200) if not chunks: logger.debug("No chunks generated for file %s", file_id) return # Delete existing chunks for this file (idempotent re-index) await db.execute( text("DELETE FROM document_chunks WHERE file_id = :fid"), {"fid": eid}, ) # Generate embeddings and insert chunks in batches for chunk in chunks: try: embedding = await generate_embedding( chunk["chunk_text"], db=db, tenant_id=tenant_id ) if embedding: await db.execute( text( "INSERT INTO document_chunks " "(tenant_id, file_id, chunk_index, chunk_text, chunk_hash, embedding) " "VALUES (:tid, :fid, :idx, :ctext, :chash, cast(:emb AS vector))" ), { "tid": tenant_id, "fid": eid, "idx": chunk["chunk_index"], "ctext": chunk["chunk_text"], "chash": chunk["chunk_hash"], "emb": str(embedding), }, ) else: logger.warning("Empty embedding for chunk %d of file %s", chunk["chunk_index"], file_id) except Exception: logger.exception("Failed to embed chunk %d for file %s", chunk["chunk_index"], file_id) await db.commit() logger.info("Indexed %d chunks for file %s", len(chunks), file_id) except Exception: logger.exception("Failed to index file chunks for %s", file_id) try: await db.rollback() except Exception: pass async def reindex_chunks(ctx: dict[str, Any], file_id: str) -> None: """Re-chunk and re-embed a file — delegates to index_file_chunks.""" await index_file_chunks(ctx, file_id) # ── Register all job functions with the job registry ────────────────────────── from app.core.job_registry import register_job # noqa: E402 register_job("index_mails", index_mails) register_job("index_file", index_file) register_job("index_contact", index_contact) register_job("index_event", index_event) register_job("reindex", reindex) register_job("reindex_all", reindex_all) register_job("embedding_batch", embedding_batch) register_job("delete_entity_index", delete_entity_index) register_job("delete_file_chunks", delete_file_chunks) register_job("retry_failed_index", retry_failed_index) register_job("index_file_chunks", index_file_chunks) register_job("reindex_chunks", reindex_chunks)