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

"""Embedding pipeline using LiteLLM with OpenRouter for embeddings.

Delegates credential lookup, model building, and embedding generation to
the centralised ``app.ai.llm_client`` module. The wrapper functions here
preserve backward compatibility for existing call sites.
"""

from __future__ import annotations

import logging
import os
import uuid
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from sqlalchemy.ext.asyncio import AsyncSession

from app.ai.llm_client import (
    EMBEDDING_DIMENSIONS,
    MAX_INPUT_CHARS,
    OPENROUTER_EMBEDDING_MODEL,
    llm_embed,
)
from app.ai.llm_client import (
    build_model as _central_build_model,
)
from app.ai.llm_client import (
    get_api_credentials as _central_get_api_credentials,
)

logger = logging.getLogger(__name__)

# Re-export constants for backward compatibility
__all__ = [
    "MAX_INPUT_CHARS",
    "OPENROUTER_API_KEY",
    "OPENROUTER_BASE_URL",
    "OPENROUTER_EMBEDDING_MODEL",
    "EMBEDDING_DIMENSIONS",
    "_get_api_credentials",
    "_build_model",
    "generate_embedding",
    "generate_embeddings_batch",
    "index_entity",
]

# Re-export for backward compatibility (consumers may import these directly)
OPENROUTER_API_KEY = os.environ.get("API_KEY_OPENROUTER", "")
OPENROUTER_BASE_URL = "https://openrouter.ai/api/v1"


async def _get_api_credentials(
    db: AsyncSession | None, tenant_id: uuid.UUID | None
) -> tuple[str | None, str | None, str | None]:
    """Get API key, base_url and provider_type for embeddings.

    Thin wrapper delegating to ``app.ai.llm_client.get_api_credentials``.
    Kept for backward compatibility with existing call sites.
    """
    return await _central_get_api_credentials(db, tenant_id)


def _build_model(model: str, provider_type: str | None) -> str:
    """Build litellm model string with provider prefix.

    Thin wrapper delegating to ``app.ai.llm_client.build_model``.
    Kept for backward compatibility with existing call sites.
    """
    return _central_build_model(model, provider_type)


async def generate_embedding(
    text: str,
    model: str | None = None,
    db: AsyncSession | None = None,
    tenant_id: uuid.UUID | None = None,
) -> list[float]:
    """Generate a single embedding via LiteLLM.

    Uses OpenRouter with text-embedding-3-small (768 dimensions).

    Args:
        text: Input text (truncated to 8000 chars).
        model: Embedding model name (default: openai/text-embedding-3-small).
        db: Optional DB session for API key lookup.
        tenant_id: Optional tenant ID for API key lookup.

    Returns:
        Embedding vector as list of floats.
    """
    embeddings = await llm_embed(
        texts=text,
        model=model,
        db=db,
        tenant_id=tenant_id,
    )
    if embeddings and embeddings[0]:
        return embeddings[0]
    return []


async def generate_embeddings_batch(
    texts: list[str],
    model: str | None = None,
    db: AsyncSession | None = None,
    tenant_id: uuid.UUID | None = None,
) -> list[list[float]]:
    """Generate embeddings for multiple texts in a single API call.

    Args:
        texts: List of input texts.
        model: Embedding model name (default: openai/text-embedding-3-small).
        db: Optional DB session for API key lookup.
        tenant_id: Optional tenant ID for API key lookup.

    Returns:
        List of embedding vectors.
    """
    return await llm_embed(
        texts=texts,
        model=model,
        db=db,
        tenant_id=tenant_id,
    )


async def index_entity(
    entity_type: str,
    entity_id: uuid.UUID,
    tenant_id: uuid.UUID,
    db: AsyncSession,
) -> bool:
    """Generate and store embedding for a single entity.

    Uses the provider registry to get embedding text, generates embedding,
    and updates the entity's embedding column.

    Returns True on success, False on failure.
    """
    from sqlalchemy import text as sql_text

    from app.plugins.builtins.unified_search.models import SearchIndexLog

    async def _log_index(action: str, status: str, error: str | None = None) -> None:
        """Write a SearchIndexLog entry for audit/debugging."""
        try:
            log_entry = SearchIndexLog(
                tenant_id=tenant_id,
                entity_type=entity_type,
                entity_id=entity_id,
                action=action,
                status=status,
                error_message=error,
            )
            db.add(log_entry)
            await db.commit()
        except Exception:
            logger.debug("Failed to write SearchIndexLog", exc_info=True)

    # ── Dedup check: skip if already indexed with unchanged content ──
    table_map = {
        "contact": "contacts",
        "mail": "mails",
        "file": "files",
        "event": "calendar_entries",
    }
    table = table_map.get(entity_type)
    if not table:
        logger.warning("Unknown entity_type=%s for embedding storage", entity_type)
        await _log_index("index", "failed", f"Unknown entity_type={entity_type}")
        return False

    # For files: compare content_hash; for others: compare updated_at vs indexed_at
    if entity_type == "file":
        dedup_result = await db.execute(
            sql_text(
                f"SELECT content_hash, indexed_at FROM {table} "
                f"WHERE id = :eid AND tenant_id = :tid"
            ),
            {"eid": entity_id, "tid": tenant_id},
        )
        dedup_row = dedup_result.mappings().first()
        if dedup_row and dedup_row.get("indexed_at") and dedup_row.get("content_hash"):
            # Already indexed and content_hash hasn't changed
            logger.debug("Skipping duplicate index for file %s (content_hash unchanged)", entity_id)
            await _log_index("index", "skipped_duplicate")
            return True
    else:
        dedup_result = await db.execute(
            sql_text(
                f"SELECT updated_at, indexed_at FROM {table} "
                f"WHERE id = :eid AND tenant_id = :tid"
            ),
            {"eid": entity_id, "tid": tenant_id},
        )
        dedup_row = dedup_result.mappings().first()
        if (
            dedup_row
            and dedup_row.get("indexed_at")
            and dedup_row.get("updated_at")
            and dedup_row["updated_at"] <= dedup_row["indexed_at"]
        ):
            logger.debug("Skipping duplicate index for %s/%s (content unchanged)", entity_type, entity_id)
            await _log_index("index", "skipped_duplicate")
            return True

    from app.plugins.builtins.unified_search.provider_registry import get_search_registry

    registry = get_search_registry()
    provider = registry.get(entity_type)
    if provider is None:
        logger.warning("No provider for entity_type=%s", entity_type)
        await _log_index("index", "failed", f"No provider for entity_type={entity_type}")
        return False

    try:
        text = await provider.get_embedding_text(db, entity_id, tenant_id)
        if not text.strip():
            logger.debug("Empty embedding text for %s/%s", entity_type, entity_id)
            await _log_index("index", "skipped_empty")
            return False

        # Apply sensitive-data filter: ensure no sensitive fields leak into
        # embedding text.  The provider builds text from DB columns, so we
        # rely on the provider selecting only non-sensitive columns.  This
        # is a secondary safety net — providers should use filter_for_embeddings
        # when constructing embedding text from dict-like data.
        from app.core.sensitive_data import get_sensitive_fields

        sensitive = get_sensitive_fields(entity_type)
        if sensitive:
            # If any sensitive field name appears as a substring in the text,
            # it's likely a key=value pair — redact it.  This is a best-effort
            # guard; providers are expected to exclude sensitive columns at
            # the SQL level.
            for field_name in sensitive:
                # Only redact if the field name appears as a key-like pattern
                import re
                text = re.sub(
                    rf"\b{re.escape(field_name)}\s*[=:]\s*\S+",
                    f"{field_name}=***REDACTED***",
                    text,
                    flags=re.IGNORECASE,
                )

        try:
            embedding = await generate_embedding(text, db=db, tenant_id=tenant_id)
            if not embedding:
                await _log_index("index", "failed", "Empty embedding returned")
                return False

            # Update the entity's embedding column + indexed_at timestamp
            sql = sql_text(
                f"UPDATE {table} SET embedding = cast(:emb AS vector), indexed_at = now() "
                f"WHERE id = :eid AND tenant_id = :tid"
            )
            await db.execute(
                sql,
                {"emb": str(embedding), "eid": entity_id, "tid": tenant_id},
            )
            await db.commit()
            await _log_index("index", "success")
            return True
        except Exception as exc:
            await db.rollback()
            await _log_index("index", "failed", str(exc))
            raise
    except Exception:
        logger.warning("Failed to index entity %s/%s", entity_type, entity_id, exc_info=True)
        return False