Skip to content

Embedding

Modules:

Name Description
base_embedding
openai_embedding

Classes:

Name Description
BaseEmbedding
OpenAIEmbedding

Attributes

__all__ module-attribute 

__all__ = ['BaseEmbedding', 'OpenAIEmbedding']

Classes

BaseEmbedding

Bases: ABC

Methods:

Name Description
aembed

Generate embedding for a single text.
aembed_batch

Generate embeddings for a list of texts.
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
dimension int

Synchronous wrapper for that runs the async implementation.
Source code in pyagenity/store/embedding/base_embedding.py 
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
class BaseEmbedding(ABC):
    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))

    @abstractmethod
    async def aembed_batch(self, texts: list[str]) -> list[list[float]]:
        """Generate embeddings for a list of texts."""
        # pragma: no cover

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

    @abstractmethod
    async def aembed(self, text: str) -> list[float]:
        """Generate embedding for a single text."""
        raise NotImplementedError

    @property
    @abstractmethod
    def dimension(self) -> int:
        """Synchronous wrapper for that runs the async implementation."""
        raise NotImplementedError

Attributes

dimension abstractmethod property 
dimension

Synchronous wrapper for that runs the async implementation.

Functions

aembed abstractmethod async 
aembed(text)

Generate embedding for a single text.

Source code in pyagenity/store/embedding/base_embedding.py 
20
21
22
23
@abstractmethod
async def aembed(self, text: str) -> list[float]:
    """Generate embedding for a single text."""
    raise NotImplementedError
aembed_batch abstractmethod async 
aembed_batch(texts)

Generate embeddings for a list of texts.

Source code in pyagenity/store/embedding/base_embedding.py 
11
12
13
@abstractmethod
async def aembed_batch(self, texts: list[str]) -> list[list[float]]:
    """Generate embeddings for a list of texts."""
embed 
embed(text)

Synchronous wrapper for aembed that runs the async implementation.

Source code in pyagenity/store/embedding/base_embedding.py 
16
17
18
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 
7
8
9
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))

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 
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
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 
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
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 
54
55
56
57
58
59
60
61
62
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 
44
45
46
47
48
49
50
51
52
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 
16
17
18
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 
7
8
9
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))

Modules

base_embedding

Classes:

Name Description
BaseEmbedding

Classes

BaseEmbedding

Bases: ABC

Methods:

Name Description
aembed

Generate embedding for a single text.
aembed_batch

Generate embeddings for a list of texts.
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
dimension int

Synchronous wrapper for that runs the async implementation.
Source code in pyagenity/store/embedding/base_embedding.py 
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
class BaseEmbedding(ABC):
    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))

    @abstractmethod
    async def aembed_batch(self, texts: list[str]) -> list[list[float]]:
        """Generate embeddings for a list of texts."""
        # pragma: no cover

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

    @abstractmethod
    async def aembed(self, text: str) -> list[float]:
        """Generate embedding for a single text."""
        raise NotImplementedError

    @property
    @abstractmethod
    def dimension(self) -> int:
        """Synchronous wrapper for that runs the async implementation."""
        raise NotImplementedError
Attributes
dimension abstractmethod property 
dimension

Synchronous wrapper for that runs the async implementation.

Functions
aembed abstractmethod async 
aembed(text)

Generate embedding for a single text.

Source code in pyagenity/store/embedding/base_embedding.py 
20
21
22
23
@abstractmethod
async def aembed(self, text: str) -> list[float]:
    """Generate embedding for a single text."""
    raise NotImplementedError
aembed_batch abstractmethod async 
aembed_batch(texts)

Generate embeddings for a list of texts.

Source code in pyagenity/store/embedding/base_embedding.py 
11
12
13
@abstractmethod
async def aembed_batch(self, texts: list[str]) -> list[list[float]]:
    """Generate embeddings for a list of texts."""
embed 
embed(text)

Synchronous wrapper for aembed that runs the async implementation.

Source code in pyagenity/store/embedding/base_embedding.py 
16
17
18
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 
7
8
9
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))

Functions

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 
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
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 
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
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 
54
55
56
57
58
59
60
61
62
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 
44
45
46
47
48
49
50
51
52
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 
16
17
18
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 
7
8
9
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))