danswer/backend/model_server/encoders.py

import gc
import json
from typing import Any
from typing import Optional

import openai
import vertexai  # type: ignore
import voyageai  # type: ignore
from cohere import Client as CohereClient
from fastapi import APIRouter
from fastapi import HTTPException
from google.oauth2 import service_account  # type: ignore
from retry import retry
from sentence_transformers import CrossEncoder  # type: ignore
from sentence_transformers import SentenceTransformer  # type: ignore
from vertexai.language_models import TextEmbeddingInput  # type: ignore
from vertexai.language_models import TextEmbeddingModel  # type: ignore

from danswer.utils.logger import setup_logger
from model_server.constants import DEFAULT_COHERE_MODEL
from model_server.constants import DEFAULT_OPENAI_MODEL
from model_server.constants import DEFAULT_VERTEX_MODEL
from model_server.constants import DEFAULT_VOYAGE_MODEL
from model_server.constants import EmbeddingModelTextType
from model_server.constants import EmbeddingProvider
from model_server.constants import MODEL_WARM_UP_STRING
from model_server.utils import simple_log_function_time
from shared_configs.configs import CROSS_EMBED_CONTEXT_SIZE
from shared_configs.configs import CROSS_ENCODER_MODEL_ENSEMBLE
from shared_configs.configs import INDEXING_ONLY
from shared_configs.enums import EmbedTextType
from shared_configs.model_server_models import EmbedRequest
from shared_configs.model_server_models import EmbedResponse
from shared_configs.model_server_models import RerankRequest
from shared_configs.model_server_models import RerankResponse


logger = setup_logger()

router = APIRouter(prefix="/encoder")

_GLOBAL_MODELS_DICT: dict[str, "SentenceTransformer"] = {}
_RERANK_MODELS: Optional[list["CrossEncoder"]] = None
# If we are not only indexing, dont want retry very long
_RETRY_DELAY = 10 if INDEXING_ONLY else 0.1
_RETRY_TRIES = 10 if INDEXING_ONLY else 2


def _initialize_client(
    api_key: str, provider: EmbeddingProvider, model: str | None = None
) -> Any:
    if provider == EmbeddingProvider.OPENAI:
        return openai.OpenAI(api_key=api_key)
    elif provider == EmbeddingProvider.COHERE:
        return CohereClient(api_key=api_key)
    elif provider == EmbeddingProvider.VOYAGE:
        return voyageai.Client(api_key=api_key)
    elif provider == EmbeddingProvider.GOOGLE:
        credentials = service_account.Credentials.from_service_account_info(
            json.loads(api_key)
        )
        project_id = json.loads(api_key)["project_id"]
        vertexai.init(project=project_id, credentials=credentials)
        return TextEmbeddingModel.from_pretrained(model or DEFAULT_VERTEX_MODEL)
    else:
        raise ValueError(f"Unsupported provider: {provider}")


class CloudEmbedding:
    def __init__(
        self,
        api_key: str,
        provider: str,
        # Only for Google as is needed on client setup
        model: str | None = None,
    ) -> None:
        try:
            self.provider = EmbeddingProvider(provider.lower())
        except ValueError:
            raise ValueError(f"Unsupported provider: {provider}")
        self.client = _initialize_client(api_key, self.provider, model)

    def _embed_openai(self, texts: list[str], model: str | None) -> list[list[float]]:
        if model is None:
            model = DEFAULT_OPENAI_MODEL

        # OpenAI does not seem to provide truncation option, however
        # the context lengths used by Danswer currently are smaller than the max token length
        # for OpenAI embeddings so it's not a big deal
        response = self.client.embeddings.create(input=texts, model=model)
        return [embedding.embedding for embedding in response.data]

    def _embed_cohere(
        self, texts: list[str], model: str | None, embedding_type: str
    ) -> list[list[float]]:
        if model is None:
            model = DEFAULT_COHERE_MODEL

        # Does not use the same tokenizer as the Danswer API server but it's approximately the same
        # empirically it's only off by a very few tokens so it's not a big deal
        response = self.client.embed(
            texts=texts,
            model=model,
            input_type=embedding_type,
            truncate="END",
        )
        return response.embeddings

    def _embed_voyage(
        self, texts: list[str], model: str | None, embedding_type: str
    ) -> list[list[float]]:
        if model is None:
            model = DEFAULT_VOYAGE_MODEL

        # Similar to Cohere, the API server will do approximate size chunking
        # it's acceptable to miss by a few tokens
        response = self.client.embed(
            texts,
            model=model,
            input_type=embedding_type,
            truncation=True,  # Also this is default
        )
        return response.embeddings

    def _embed_vertex(
        self, texts: list[str], model: str | None, embedding_type: str
    ) -> list[list[float]]:
        if model is None:
            model = DEFAULT_VERTEX_MODEL

        embeddings = self.client.get_embeddings(
            [
                TextEmbeddingInput(
                    text,
                    embedding_type,
                )
                for text in texts
            ],
            auto_truncate=True,  # Also this is default
        )
        return [embedding.values for embedding in embeddings]

    @retry(tries=_RETRY_TRIES, delay=_RETRY_DELAY)
    def embed(
        self,
        *,
        texts: list[str],
        text_type: EmbedTextType,
        model_name: str | None = None,
    ) -> list[list[float]]:
        try:
            if self.provider == EmbeddingProvider.OPENAI:
                return self._embed_openai(texts, model_name)

            embedding_type = EmbeddingModelTextType.get_type(self.provider, text_type)
            if self.provider == EmbeddingProvider.COHERE:
                return self._embed_cohere(texts, model_name, embedding_type)
            elif self.provider == EmbeddingProvider.VOYAGE:
                return self._embed_voyage(texts, model_name, embedding_type)
            elif self.provider == EmbeddingProvider.GOOGLE:
                return self._embed_vertex(texts, model_name, embedding_type)
            else:
                raise ValueError(f"Unsupported provider: {self.provider}")
        except Exception as e:
            raise HTTPException(
                status_code=500,
                detail=f"Error embedding text with {self.provider}: {str(e)}",
            )

    @staticmethod
    def create(
        api_key: str, provider: str, model: str | None = None
    ) -> "CloudEmbedding":
        logger.debug(f"Creating Embedding instance for provider: {provider}")
        return CloudEmbedding(api_key, provider, model)


def get_embedding_model(
    model_name: str,
    max_context_length: int,
) -> "SentenceTransformer":
    from sentence_transformers import SentenceTransformer  # type: ignore

    global _GLOBAL_MODELS_DICT  # A dictionary to store models

    if _GLOBAL_MODELS_DICT is None:
        _GLOBAL_MODELS_DICT = {}

    if model_name not in _GLOBAL_MODELS_DICT:
        logger.info(f"Loading {model_name}")
        model = SentenceTransformer(model_name)
        model.max_seq_length = max_context_length
        _GLOBAL_MODELS_DICT[model_name] = model
    elif max_context_length != _GLOBAL_MODELS_DICT[model_name].max_seq_length:
        _GLOBAL_MODELS_DICT[model_name].max_seq_length = max_context_length

    return _GLOBAL_MODELS_DICT[model_name]


def get_local_reranking_model_ensemble(
    model_names: list[str] = CROSS_ENCODER_MODEL_ENSEMBLE,
    max_context_length: int = CROSS_EMBED_CONTEXT_SIZE,
) -> list[CrossEncoder]:
    global _RERANK_MODELS
    if _RERANK_MODELS is None or max_context_length != _RERANK_MODELS[0].max_length:
        del _RERANK_MODELS
        gc.collect()

        _RERANK_MODELS = []
        for model_name in model_names:
            logger.info(f"Loading {model_name}")
            model = CrossEncoder(model_name)
            model.max_length = max_context_length
            _RERANK_MODELS.append(model)
    return _RERANK_MODELS


def warm_up_cross_encoders() -> None:
    logger.info(f"Warming up Cross-Encoders: {CROSS_ENCODER_MODEL_ENSEMBLE}")

    cross_encoders = get_local_reranking_model_ensemble()
    [
        cross_encoder.predict((MODEL_WARM_UP_STRING, MODEL_WARM_UP_STRING))
        for cross_encoder in cross_encoders
    ]


@simple_log_function_time()
def embed_text(
    texts: list[str],
    text_type: EmbedTextType,
    model_name: str | None,
    max_context_length: int,
    normalize_embeddings: bool,
    api_key: str | None,
    provider_type: str | None,
    prefix: str | None,
) -> list[list[float]]:
    # Third party API based embedding model
    if provider_type is not None:
        logger.debug(f"Embedding text with provider: {provider_type}")
        if api_key is None:
            raise RuntimeError("API key not provided for cloud model")

        if prefix:
            # This may change in the future if some providers require the user
            # to manually append a prefix but this is not the case currently
            raise ValueError(
                "Prefix string is not valid for cloud models. "
                "Cloud models take an explicit text type instead."
            )

        cloud_model = CloudEmbedding(
            api_key=api_key, provider=provider_type, model=model_name
        )
        embeddings = cloud_model.embed(
            texts=texts,
            model_name=model_name,
            text_type=text_type,
        )

    # Locally running model
    elif model_name is not None:
        prefixed_texts = [f"{prefix}{text}" for text in texts] if prefix else texts
        local_model = get_embedding_model(
            model_name=model_name, max_context_length=max_context_length
        )
        embeddings = local_model.encode(
            prefixed_texts, normalize_embeddings=normalize_embeddings
        )

    else:
        raise ValueError(
            "Either model name or provider must be provided to run embeddings."
        )

    if embeddings is None:
        raise RuntimeError("Failed to create Embeddings")

    if not isinstance(embeddings, list):
        embeddings = embeddings.tolist()

    return embeddings


@simple_log_function_time()
def calc_sim_scores(query: str, docs: list[str]) -> list[list[float]]:
    cross_encoders = get_local_reranking_model_ensemble()
    sim_scores = [
        encoder.predict([(query, doc) for doc in docs]).tolist()  # type: ignore
        for encoder in cross_encoders
    ]
    return sim_scores


@router.post("/bi-encoder-embed")
async def process_embed_request(
    embed_request: EmbedRequest,
) -> EmbedResponse:
    if not embed_request.texts:
        raise HTTPException(status_code=400, detail="No texts to be embedded")

    try:
        if embed_request.text_type == EmbedTextType.QUERY:
            prefix = embed_request.manual_query_prefix
        elif embed_request.text_type == EmbedTextType.PASSAGE:
            prefix = embed_request.manual_passage_prefix
        else:
            prefix = None

        embeddings = embed_text(
            texts=embed_request.texts,
            model_name=embed_request.model_name,
            max_context_length=embed_request.max_context_length,
            normalize_embeddings=embed_request.normalize_embeddings,
            api_key=embed_request.api_key,
            provider_type=embed_request.provider_type,
            text_type=embed_request.text_type,
            prefix=prefix,
        )
        return EmbedResponse(embeddings=embeddings)
    except Exception as e:
        exception_detail = f"Error during embedding process:\n{str(e)}"
        logger.exception(exception_detail)
        raise HTTPException(status_code=500, detail=exception_detail)


@router.post("/cross-encoder-scores")
async def process_rerank_request(embed_request: RerankRequest) -> RerankResponse:
    """Cross encoders can be purely black box from the app perspective"""
    if INDEXING_ONLY:
        raise RuntimeError("Indexing model server should not call intent endpoint")

    if not embed_request.documents or not embed_request.query:
        raise HTTPException(
            status_code=400, detail="No documents or query to be reranked"
        )

    try:
        sim_scores = calc_sim_scores(
            query=embed_request.query, docs=embed_request.documents
        )
        return RerankResponse(scores=sim_scores)
    except Exception as e:
        logger.exception(f"Error during reranking process:\n{str(e)}")
        raise HTTPException(
            status_code=500, detail="Failed to run Cross-Encoder reranking"
        )