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,
    build_model as _central_build_model,
    get_api_credentials as _central_get_api_credentials,
    llm_embed,
)

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 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)
        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)
            return False

        embedding = await generate_embedding(text, db=db, tenant_id=tenant_id)
        if not embedding:
            return False

        # Update the entity's embedding column
        from sqlalchemy import text as sql_text

        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)
            return False

        sql = sql_text(
            f"UPDATE {table} SET embedding = cast(:emb AS vector) "
            f"WHERE id = :eid AND tenant_id = :tid"
        )
        await db.execute(
            sql,
            {"emb": str(embedding), "eid": entity_id, "tid": tenant_id},
        )
        await db.commit()
        return True
    except Exception:
        logger.warning("Failed to index entity %s/%s", entity_type, entity_id, exc_info=True)
        return False