Source code for pyrit.embedding.openai_text_embedding

# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.

from openai import OpenAI

from pyrit.embedding._text_embedding import _TextEmbedding




class OpenAiTextEmbedding(_TextEmbedding):
    """
    Provides text embedding functionality using OpenAI services.
    """



    def __init__(self, *, model: str, api_key: str) -> None:
        """
        Generate embedding using OpenAI API.

        Args:
            model: The model to use
            api_key: The API key to use
        """
        self._client = OpenAI(
            api_key=api_key,
        )
        self._model = model
        super().__init__()