"""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