danswer/backend/model_server/encoders.py

import json
from typing import Any
from typing import Optional

import httpx
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.utils import simple_log_function_time
from shared_configs.configs import INDEXING_ONLY
from shared_configs.enums import EmbedTextType
from shared_configs.enums import RerankerProvider
from shared_configs.model_server_models import Embedding
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
from shared_configs.utils import batch_list


logger = setup_logger()

router = APIRouter(prefix="/encoder")

_GLOBAL_MODELS_DICT: dict[str, "SentenceTransformer"] = {}
_RERANK_MODEL: Optional["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

# OpenAI only allows 2048 embeddings to be computed at once
_OPENAI_MAX_INPUT_LEN = 2048
# Cohere allows up to 96 embeddings in a single embedding calling
_COHERE_MAX_INPUT_LEN = 96


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: EmbeddingProvider,
        # Only for Google as is needed on client setup
        model: str | None = None,
    ) -> None:
        self.provider = provider
        self.client = _initialize_client(api_key, self.provider, model)

    def _embed_openai(self, texts: list[str], model: str | None) -> list[Embedding]:
        if not model:
            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
        final_embeddings: list[Embedding] = []
        try:
            for text_batch in batch_list(texts, _OPENAI_MAX_INPUT_LEN):
                response = self.client.embeddings.create(input=text_batch, model=model)
                final_embeddings.extend(
                    [embedding.embedding for embedding in response.data]
                )
            return final_embeddings
        except Exception as e:
            error_string = (
                f"Error embedding text with OpenAI: {str(e)} \n"
                f"Model: {model} \n"
                f"Provider: {self.provider} \n"
                f"Texts: {texts}"
            )
            logger.error(error_string)
            raise RuntimeError(error_string)

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

        final_embeddings: list[Embedding] = []
        for text_batch in batch_list(texts, _COHERE_MAX_INPUT_LEN):
            # 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=text_batch,
                model=model,
                input_type=embedding_type,
                truncate="END",
            )
            final_embeddings.extend(response.embeddings)
        return final_embeddings

    def _embed_voyage(
        self, texts: list[str], model: str | None, embedding_type: str
    ) -> list[Embedding]:
        if not model:
            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[Embedding]:
        if not model:
            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[Embedding]:
        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: EmbeddingProvider, 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.notice(f"Loading {model_name}")
        # Some model architectures that aren't built into the Transformers or Sentence
        # Transformer need to be downloaded to be loaded locally. This does not mean
        # data is sent to remote servers for inference, however the remote code can
        # be fairly arbitrary so only use trusted models
        model = SentenceTransformer(
            model_name_or_path=model_name,
            trust_remote_code=True,
        )
        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(
    model_name: str,
) -> CrossEncoder:
    global _RERANK_MODEL
    if _RERANK_MODEL is None:
        logger.notice(f"Loading {model_name}")
        model = CrossEncoder(model_name)
        _RERANK_MODEL = model
    return _RERANK_MODEL


def embed_with_litellm_proxy(
    texts: list[str], api_url: str, model_name: str, api_key: str | None
) -> list[Embedding]:
    headers = {} if not api_key else {"Authorization": f"Bearer {api_key}"}

    with httpx.Client() as client:
        response = client.post(
            api_url,
            json={
                "model": model_name,
                "input": texts,
            },
            headers=headers,
        )
        response.raise_for_status()
        result = response.json()
        return [embedding["embedding"] for embedding in result["data"]]


@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: EmbeddingProvider | None,
    prefix: str | None,
    api_url: str | None,
) -> list[Embedding]:
    logger.info(f"Embedding {len(texts)} texts with provider: {provider_type}")

    if not all(texts):
        logger.error("Empty strings provided for embedding")
        raise ValueError("Empty strings are not allowed for embedding.")

    if not texts:
        logger.error("No texts provided for embedding")
        raise ValueError("No texts provided for embedding.")

    if provider_type == EmbeddingProvider.LITELLM:
        logger.debug(f"Using LiteLLM proxy for embedding with URL: {api_url}")
        if not api_url:
            logger.error("API URL not provided for LiteLLM proxy")
            raise ValueError("API URL is required for LiteLLM proxy embedding.")
        try:
            return embed_with_litellm_proxy(
                texts=texts,
                api_url=api_url,
                model_name=model_name or "",
                api_key=api_key,
            )
        except Exception as e:
            logger.exception(f"Error during LiteLLM proxy embedding: {str(e)}")
            raise

    elif provider_type is not None:
        logger.debug(f"Using cloud provider {provider_type} for embedding")
        if api_key is None:
            logger.error("API key not provided for cloud model")
            raise RuntimeError("API key not provided for cloud model")

        if prefix:
            logger.warning("Prefix provided for cloud model, which is not supported")
            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,
        )

        if any(embedding is None for embedding in embeddings):
            error_message = "Embeddings contain None values\n"
            error_message += "Corresponding texts:\n"
            error_message += "\n".join(texts)
            logger.error(error_message)
            raise ValueError(error_message)

    elif model_name is not None:
        logger.debug(f"Using local model {model_name} for embedding")
        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_vectors = local_model.encode(
            prefixed_texts, normalize_embeddings=normalize_embeddings
        )
        embeddings = [
            embedding if isinstance(embedding, list) else embedding.tolist()
            for embedding in embeddings_vectors
        ]

    else:
        logger.error("Neither model name nor provider specified for embedding")
        raise ValueError(
            "Either model name or provider must be provided to run embeddings."
        )

    logger.info(f"Successfully embedded {len(texts)} texts")
    return embeddings


@simple_log_function_time()
def local_rerank(query: str, docs: list[str], model_name: str) -> list[float]:
    cross_encoder = get_local_reranking_model(model_name)
    return cross_encoder.predict([(query, doc) for doc in docs]).tolist()  # type: ignore


def cohere_rerank(
    query: str, docs: list[str], model_name: str, api_key: str
) -> list[float]:
    cohere_client = CohereClient(api_key=api_key)
    response = cohere_client.rerank(query=query, documents=docs, model=model_name)
    results = response.results
    sorted_results = sorted(results, key=lambda item: item.index)
    return [result.relevance_score for result in sorted_results]


def litellm_rerank(
    query: str, docs: list[str], api_url: str, model_name: str, api_key: str | None
) -> list[float]:
    headers = {} if not api_key else {"Authorization": f"Bearer {api_key}"}
    with httpx.Client() as client:
        response = client.post(
            api_url,
            json={
                "model": model_name,
                "query": query,
                "documents": docs,
            },
            headers=headers,
        )
        response.raise_for_status()
        result = response.json()
        return [
            item["relevance_score"]
            for item in sorted(result["results"], key=lambda x: x["index"])
        ]


@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")
    elif not all(embed_request.texts):
        raise ValueError("Empty strings are not allowed for embedding.")

    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,
            api_url=embed_request.api_url,
            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(rerank_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 rerank_request.documents or not rerank_request.query:
        raise HTTPException(
            status_code=400, detail="Missing documents or query for reranking"
        )
    if not all(rerank_request.documents):
        raise ValueError("Empty documents cannot be reranked.")

    try:
        if rerank_request.provider_type is None:
            sim_scores = local_rerank(
                query=rerank_request.query,
                docs=rerank_request.documents,
                model_name=rerank_request.model_name,
            )
            return RerankResponse(scores=sim_scores)
        elif rerank_request.provider_type == RerankerProvider.LITELLM:
            if rerank_request.api_url is None:
                raise ValueError("API URL is required for LiteLLM reranking.")

            sim_scores = litellm_rerank(
                query=rerank_request.query,
                docs=rerank_request.documents,
                api_url=rerank_request.api_url,
                model_name=rerank_request.model_name,
                api_key=rerank_request.api_key,
            )

            return RerankResponse(scores=sim_scores)

        elif rerank_request.provider_type == RerankerProvider.COHERE:
            if rerank_request.api_key is None:
                raise RuntimeError("Cohere Rerank Requires an API Key")
            sim_scores = cohere_rerank(
                query=rerank_request.query,
                docs=rerank_request.documents,
                model_name=rerank_request.model_name,
                api_key=rerank_request.api_key,
            )
            return RerankResponse(scores=sim_scores)
        else:
            raise ValueError(f"Unsupported provider: {rerank_request.provider_type}")
    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"
        )