Agent Zero
dc24c37c19
- Ollama Cloud has no embedding endpoint, OpenRouter does - text-embedding-3-small with dimensions=768 matches DB vector(768) column - API_KEY_OPENROUTER env var is primary, DB provider is fallback - Added dimensions parameter for text-embedding-3 models
207 lines
6.8 KiB
Python
207 lines
6.8 KiB
Python
"""Embedding pipeline using LiteLLM with OpenRouter for embeddings."""
|
||
|
||
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
|
||
|
||
# OpenRouter for embeddings (Ollama Cloud has no embedding endpoint)
|
||
OPENROUTER_API_KEY = os.environ.get('API_KEY_OPENROUTER', '')
|
||
OPENROUTER_BASE_URL = 'https://openrouter.ai/api/v1'
|
||
OPENROUTER_EMBEDDING_MODEL = os.environ.get('SEARCH_EMBEDDING_MODEL', 'openai/text-embedding-3-small')
|
||
EMBEDDING_DIMENSIONS = 768 # Must match DB column vector(768)
|
||
|
||
|
||
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.
|
||
|
||
Priority:
|
||
1. OpenRouter env var (API_KEY_OPENROUTER) – dedicated embedding provider
|
||
2. Default AI provider from DB (fallback)
|
||
3. API_KEY_OLLAMA_CLOUD env var (last resort)
|
||
"""
|
||
# OpenRouter is the primary embedding provider
|
||
if OPENROUTER_API_KEY:
|
||
return OPENROUTER_API_KEY, OPENROUTER_BASE_URL, 'openai'
|
||
|
||
# Fallback to DB provider
|
||
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.
|
||
|
||
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.
|
||
"""
|
||
model = model or OPENROUTER_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
|
||
|
||
# Request 768 dimensions to match DB vector(768) column
|
||
if 'text-embedding-3' in litellm_model:
|
||
litellm_kwargs['dimensions'] = EMBEDDING_DIMENSIONS
|
||
|
||
response = await litellm.aembedding(**litellm_kwargs)
|
||
return response.data[0]["embedding"]
|
||
except Exception:
|
||
logger.warning("Failed to generate embedding", 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: 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.
|
||
"""
|
||
model = model or OPENROUTER_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
|
||
|
||
if 'text-embedding-3' in litellm_model:
|
||
litellm_kwargs['dimensions'] = EMBEDDING_DIMENSIONS
|
||
|
||
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
|