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

"""Embedding pipeline using LiteLLM with configurable Ollama Cloud provider."""

from __future__ import annotations

import os
import logging
import uuid
from typing import Any, TYPE_CHECKING

import litellm

if TYPE_CHECKING:
    from sqlalchemy.ext.asyncio import AsyncSession

logger = logging.getLogger(__name__)

MAX_INPUT_CHARS = 8000

DEFAULT_EMBEDDING_MODEL = os.environ.get('SEARCH_EMBEDDING_MODEL', 'ollama/nomic-embed-text')


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 from the default AI provider in DB.

    Falls back to API_KEY_OLLAMA_CLOUD env var.
    Returns (api_key, base_url, provider_type).
    """
    if db and tenant_id:
        try:
            from app.plugins.builtins.ai_assistant.services import get_default_provider
            provider = await get_default_provider(db, tenant_id)
            if provider and provider.api_key:
                return provider.api_key, provider.base_url, provider.provider_type
        except Exception:
            logger.debug("Failed to get provider from DB, falling back to env")
    env_key = os.environ.get('API_KEY_OLLAMA_CLOUD', '')
    return (env_key if env_key else None), None, None


def _build_model(model: str, provider_type: str | None) -> str:
    """Build litellm model string with provider prefix."""
    if provider_type:
        model_parts = model.split("/", 1)
        return f"{provider_type}/{model_parts[-1]}"
    return model


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.

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

    Returns:
        Embedding vector as list of floats.
    """
    model = model or DEFAULT_EMBEDDING_MODEL
    truncated = text[:MAX_INPUT_CHARS]
    try:
        api_key, api_base, provider_type = await _get_api_credentials(db, tenant_id)
        litellm_model = _build_model(model, provider_type)

        litellm_kwargs: dict[str, Any] = dict(
            model=litellm_model,
            input=truncated,
        )
        if api_key:
            litellm_kwargs["api_key"] = api_key
        if api_base:
            litellm_kwargs["api_base"] = api_base

        response = await litellm.aembedding(**litellm_kwargs)
        return response.data[0]["embedding"]
    except Exception:
        logger.warning("Failed to generate embedding (Ollama Cloud may not support embeddings)", exc_info=True)
        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: ollama/nomic-embed-text).
        db: Optional DB session for API key lookup.
        tenant_id: Optional tenant ID for API key lookup.

    Returns:
        List of embedding vectors.
    """
    model = model or DEFAULT_EMBEDDING_MODEL
    truncated = [t[:MAX_INPUT_CHARS] for t in texts]
    try:
        api_key, api_base, provider_type = await _get_api_credentials(db, tenant_id)
        litellm_model = _build_model(model, provider_type)

        litellm_kwargs: dict[str, Any] = dict(
            model=litellm_model,
            input=truncated,
        )
        if api_key:
            litellm_kwargs["api_key"] = api_key
        if api_base:
            litellm_kwargs["api_base"] = api_base

        response = await litellm.aembedding(**litellm_kwargs)
        return [d["embedding"] for d in response.data]
    except Exception:
        logger.warning("Failed to generate batch embeddings", exc_info=True)
        return [[] for _ in texts]


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",
            "company": "companies",
            "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.exception("Failed to index entity %s/%s", entity_type, entity_id)
        return False