"""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 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) 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.text_extraction import extract_text_from_file from app.plugins.builtins.unified_search.embedding import generate_embedding 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) 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) async def index_company(ctx: dict[str, Any], company_id: str) -> None: """Index a company: 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(company_id) result = await db.execute( text("SELECT tenant_id FROM companies WHERE id = :cid"), {"cid": eid}, ) row = result.mappings().first() if not row: return await index_entity("company", eid, row["tenant_id"], db) except Exception: logger.exception("Failed to index company %s", company_id) 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) async def reindex(ctx: dict[str, Any], entity_type: str) -> None: """Reindex all entities of a given type with pagination.""" from sqlalchemy import text from app.plugins.builtins.unified_search.embedding import index_entity table_map = { "contact": "contacts", "company": "companies", "mail": "mails", "file": "files", "event": "calendar_entries", } 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: 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, ) except Exception: logger.exception("Reindex failed for %s/%s", entity_type, row["id"]) offset += BATCH_SIZE logger.info("Reindex complete for %s", entity_type) 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 table_map = { "contact": "contacts", "company": "companies", "mail": "mails", "file": "files", "event": "calendar_entries", } 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"]) except Exception: logger.exception("Batch query failed for %s", etype) logger.info("Embedding batch job complete")