leocrm/app/plugins/builtins/unified_search/jobs.py

"""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 contacts 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": "contacts",
        "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": "contacts",
        "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")