Openai embedding

Classes:

Name	Description
OpenAIEmbedding

Attributes:

Name	Type	Description
HAS_OPENAI

Attributes

HAS_OPENAI module-attribute

HAS_OPENAI = True

Classes

OpenAIEmbedding

Bases: BaseEmbedding

Methods:

Name	Description
__init__
aembed
aembed_batch
embed	Synchronous wrapper for aembed that runs the async implementation.
embed_batch	Synchronous wrapper for aembed_batch that runs the async implementation.

Attributes:

Name	Type	Description
api_key
client
dimension	int
model

Source code in pyagenity/store/embedding/openai_embedding.py

class OpenAIEmbedding(BaseEmbedding):
    def __init__(
        self,
        model: str = "text-embedding-3-small",
        OPENAI_API_KEY: str | None = None,
    ) -> None:
        if not HAS_OPENAI:
            raise ImportError(
                "The 'openai' package is required for OpenAIEmbedding. "
                "Please install it via 'pip install openai'."
            )
        self.model = model
        if OPENAI_API_KEY:
            self.api_key = OPENAI_API_KEY
        elif "OPENAI_API_KEY" in os.environ:
            self.api_key = os.environ["OPENAI_API_KEY"]
        else:
            raise ValueError(
                "OpenAI API key must be provided via parameter or OPENAI_API_KEY env var"
            )

        self.client = AsyncOpenAI(
            api_key=self.api_key,
        )

    async def aembed_batch(self, texts: list[str]) -> list[list[float]]:
        try:
            response = await self.client.embeddings.create(
                input=texts,
                model=self.model,
            )
            return [data.embedding for data in response.data]
        except OpenAIError as e:
            raise RuntimeError(f"OpenAI API error: {e}") from e

    async def aembed(self, text: str) -> list[float]:
        try:
            response = await self.client.embeddings.create(
                input=text,
                model=self.model,
            )
            return response.data[0].embedding if response.data else []
        except OpenAIError as e:
            raise RuntimeError(f"OpenAI API error: {e}") from e

    @property
    def dimension(self) -> int:
        model_dimensions = {
            "text-embedding-3-small": 1536,
            "text-embedding-3-large": 1536,
            "text-embedding-3-xl": 1536,
            "text-embedding-4-base": 8192,
            "text-embedding-4-large": 8192,
        }
        if self.model in model_dimensions:
            return model_dimensions[self.model]
        raise ValueError(f"Unknown model '{self.model}'. Cannot determine dimension.")

Attributes

api_key instance-attribute

api_key = OPENAI_API_KEY

client instance-attribute

client = AsyncOpenAI(api_key=api_key)

dimension property

dimension

model instance-attribute

model = model

Functions

__init__

__init__(model='text-embedding-3-small', OPENAI_API_KEY=None)

Source code in pyagenity/store/embedding/openai_embedding.py

def __init__(
    self,
    model: str = "text-embedding-3-small",
    OPENAI_API_KEY: str | None = None,
) -> None:
    if not HAS_OPENAI:
        raise ImportError(
            "The 'openai' package is required for OpenAIEmbedding. "
            "Please install it via 'pip install openai'."
        )
    self.model = model
    if OPENAI_API_KEY:
        self.api_key = OPENAI_API_KEY
    elif "OPENAI_API_KEY" in os.environ:
        self.api_key = os.environ["OPENAI_API_KEY"]
    else:
        raise ValueError(
            "OpenAI API key must be provided via parameter or OPENAI_API_KEY env var"
        )

    self.client = AsyncOpenAI(
        api_key=self.api_key,
    )

aembed async

aembed(text)

Source code in pyagenity/store/embedding/openai_embedding.py

async def aembed(self, text: str) -> list[float]:
    try:
        response = await self.client.embeddings.create(
            input=text,
            model=self.model,
        )
        return response.data[0].embedding if response.data else []
    except OpenAIError as e:
        raise RuntimeError(f"OpenAI API error: {e}") from e

aembed_batch async

aembed_batch(texts)

Source code in pyagenity/store/embedding/openai_embedding.py

async def aembed_batch(self, texts: list[str]) -> list[list[float]]:
    try:
        response = await self.client.embeddings.create(
            input=texts,
            model=self.model,
        )
        return [data.embedding for data in response.data]
    except OpenAIError as e:
        raise RuntimeError(f"OpenAI API error: {e}") from e

embed

embed(text)

Synchronous wrapper for aembed that runs the async implementation.

Source code in pyagenity/store/embedding/base_embedding.py

def embed(self, text: str) -> list[float]:
    """Synchronous wrapper for `aembed` that runs the async implementation."""
    return run_coroutine(self.aembed(text))

embed_batch

embed_batch(texts)

Synchronous wrapper for aembed_batch that runs the async implementation.

Source code in pyagenity/store/embedding/base_embedding.py

def embed_batch(self, texts: list[str]) -> list[list[float]]:
    """Synchronous wrapper for `aembed_batch` that runs the async implementation."""
    return run_coroutine(self.aembed_batch(texts))