LiteLLMEmbeddingClient

dbally.embeddings.LiteLLMEmbeddingClient

LiteLLMEmbeddingClient(model: str = 'text-embedding-3-small', options: Optional[Dict] = None, api_base: Optional[str] = None, api_key: Optional[str] = None, api_version: Optional[str] = None)

Bases: EmbeddingClient

Client for creating text embeddings using LiteLLM API.

Constructs the LiteLLMEmbeddingClient.

PARAMETER	DESCRIPTION
model	Name of the LiteLLM supported model to be used. Default is "text-embedding-3-small". TYPE: str DEFAULT: 'text-embedding-3-small'
options	Additional options to pass to the LiteLLM API. TYPE: Optional[Dict] DEFAULT: None
api_base	The API endpoint you want to call the model with. TYPE: Optional[str] DEFAULT: None
api_key	API key to be used. API key to be used. If not specified, an environment variable will be used, for more information, follow the instructions for your specific vendor in the LiteLLM documentation. TYPE: Optional[str] DEFAULT: None
api_version	The API version for the call. TYPE: Optional[str] DEFAULT: None

RAISES	DESCRIPTION
ImportError	If the litellm package is not installed.

Source code in src/dbally/embeddings/litellm.py

def __init__(
    self,
    model: str = "text-embedding-3-small",
    options: Optional[Dict] = None,
    api_base: Optional[str] = None,
    api_key: Optional[str] = None,
    api_version: Optional[str] = None,
) -> None:
    """
    Constructs the LiteLLMEmbeddingClient.

    Args:
        model: Name of the [LiteLLM supported model](https://docs.litellm.ai/docs/embedding/supported_embedding)\
            to be used. Default is "text-embedding-3-small".
        options: Additional options to pass to the LiteLLM API.
        api_base: The API endpoint you want to call the model with.
        api_key: API key to be used. API key to be used. If not specified, an environment variable will be used,
            for more information, follow the instructions for your specific vendor in the\
            [LiteLLM documentation](https://docs.litellm.ai/docs/embedding/supported_embedding).
        api_version: The API version for the call.

    Raises:
        ImportError: If the litellm package is not installed.
    """
    if not HAVE_LITELLM:
        raise ImportError("You need to install litellm package to use LiteLLM models")

    super().__init__()
    self.model = model
    self.options = options or {}
    self.api_base = api_base
    self.api_key = api_key
    self.api_version = api_version

model instance-attribute

model = model

options instance-attribute

options = options or {}

api_base instance-attribute

api_base = api_base

api_key instance-attribute

api_key = api_key

api_version instance-attribute

api_version = api_version

get_embeddings async

get_embeddings(data: List[str]) -> List[List[float]]

Creates embeddings for the given strings.

PARAMETER	DESCRIPTION
data	List of strings to get embeddings for. TYPE: List[str]

RETURNS	DESCRIPTION
List[List[float]]	List of embeddings for the given strings.

RAISES	DESCRIPTION
EmbeddingConnectionError	If there is a connection error with the embedding API.
EmbeddingStatusError	If the embedding API returns an error status code.
EmbeddingResponseError	If the embedding API response is invalid.

Source code in src/dbally/embeddings/litellm.py

async def get_embeddings(self, data: List[str]) -> List[List[float]]:
    """
    Creates embeddings for the given strings.

    Args:
        data: List of strings to get embeddings for.

    Returns:
        List of embeddings for the given strings.

    Raises:
        EmbeddingConnectionError: If there is a connection error with the embedding API.
        EmbeddingStatusError: If the embedding API returns an error status code.
        EmbeddingResponseError: If the embedding API response is invalid.
    """
    try:
        response = await litellm.aembedding(
            input=data,
            model=self.model,
            api_base=self.api_base,
            api_key=self.api_key,
            api_version=self.api_version,
            **self.options,
        )
    except litellm.openai.APIConnectionError as exc:
        raise EmbeddingConnectionError() from exc
    except litellm.openai.APIStatusError as exc:
        raise EmbeddingStatusError(exc.message, exc.status_code) from exc
    except litellm.openai.APIResponseValidationError as exc:
        raise EmbeddingResponseError() from exc

    return [embedding["embedding"] for embedding in response.data]