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 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')
OLLAMA_API_KEY = os.environ.get('API_KEY_OLLAMA_CLOUD', '')


async def generate_embedding(text: str, model: str | 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).

    Returns:
        Embedding vector as list of floats.
    """
    model = model or DEFAULT_EMBEDDING_MODEL
    truncated = text[:MAX_INPUT_CHARS]
    try:
        response = await litellm.aembedding(
            model=model,
            input=truncated,
            api_key=OLLAMA_API_KEY,
        )
        return response.data[0]["embedding"]
    except Exception:
        logger.exception("Failed to generate embedding")
        return []


async def generate_embeddings_batch(
    texts: list[str], model: str | 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).

    Returns:
        List of embedding vectors.
    """
    model = model or DEFAULT_EMBEDDING_MODEL
    truncated = [t[:MAX_INPUT_CHARS] for t in texts]
    try:
        response = await litellm.aembedding(
            model=model,
            input=truncated,
            api_key=OLLAMA_API_KEY,
        )
        return [d["embedding"] for d in response.data]
    except Exception:
        logger.exception("Failed to generate batch embeddings")
        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)
        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