pyrit.embedding.OpenAITextEmbedding

pyrit.embedding.OpenAITextEmbedding#

class OpenAITextEmbedding(*, api_key: str | Callable[[], str | Awaitable[str]] | None = None, endpoint: str | None = None, model_name: str | None = None)[source]#

Bases: EmbeddingSupport

Text embedding class that works with both Azure OpenAI and platform OpenAI endpoints. Uses the AsyncOpenAI client under the hood for both providers since they share the same API.

__init__(*, api_key: str | Callable[[], str | Awaitable[str]] | None = None, endpoint: str | None = None, model_name: str | None = None) → None[source]#

Initialize text embedding client for Azure OpenAI or platform OpenAI.

Parameters:

  • api_key – The API key (string) or token provider (callable) for authentication. For Azure with Entra auth, pass get_azure_openai_auth(endpoint) from pyrit.auth. Defaults to OPENAI_EMBEDDING_KEY environment variable.

  • endpoint – The API endpoint URL. For Azure: https://{resource}.openai.azure.com/openai/v1 For platform OpenAI: https://api.openai.com/v1 Defaults to OPENAI_EMBEDDING_ENDPOINT environment variable.

  • model_name – The model/deployment name (e.g., “text-embedding-3-small”). Defaults to OPENAI_EMBEDDING_MODEL environment variable.

Methods

__init__(*[, api_key, endpoint, model_name])

Initialize text embedding client for Azure OpenAI or platform OpenAI.

generate_text_embedding(text, **kwargs)

Generate text embedding synchronously by calling the async method.

generate_text_embedding_async(text, **kwargs)

Generate text embedding asynchronously.

Attributes

API_KEY_ENVIRONMENT_VARIABLE

ENDPOINT_URI_ENVIRONMENT_VARIABLE

MODEL_ENVIRONMENT_VARIABLE

API_KEY_ENVIRONMENT_VARIABLE: str = 'OPENAI_EMBEDDING_KEY'#
ENDPOINT_URI_ENVIRONMENT_VARIABLE: str = 'OPENAI_EMBEDDING_ENDPOINT'#
MODEL_ENVIRONMENT_VARIABLE: str = 'OPENAI_EMBEDDING_MODEL'#
generate_text_embedding(text: str, **kwargs) → EmbeddingResponse[source]#

Generate text embedding synchronously by calling the async method.

Parameters:

  • text – The text to generate the embedding for

  • **kwargs – Additional arguments to pass to the embeddings API

Returns:

The embedding response

async generate_text_embedding_async(text: str, **kwargs) → EmbeddingResponse[source]#

Generate text embedding asynchronously.

Parameters:

  • text – The text to generate the embedding for

  • **kwargs – Additional arguments to pass to the embeddings API

Returns:

The embedding response

Contents