agentscope.embedding

The embedding module in agentscope.

class EmbeddingModelBase

Bases: object

Base class for embedding models.

__init__(model_name)

Initialize the embedding model base class.

Parameters:

model_name (str) – The name of the embedding model.

Return type:

None

model_name: str

The embedding model name

async __call__(text, **kwargs)

Call the embedding API with the given arguments.

Parameters:

• text (List[str])

• kwargs (Any)

Return type:

EmbeddingResponse

class EmbeddingUsage

Bases: DictMixin

The usage of an embedding model API invocation.

time: float

The time used in seconds.

__init__(time, tokens=<factory>, type=<factory>)

Parameters:

• time (float)

• tokens (int | None)

• type (Literal['embedding'])

Return type:

None

tokens: int | None

The number of tokens used, if available.

type: Literal['embedding']

The type of the usage, must be embedding.

class EmbeddingResponse

Bases: DictMixin

The embedding response class.

embeddings: List[List[float]]

The embedding data

id: str

The identity of the embedding response

created_at: str

The timestamp of the embedding response creation

__init__(embeddings, id=<factory>, created_at=<factory>, type=<factory>, usage=<factory>, source=<factory>)

Parameters:

• embeddings (List[List[float]])

• id (str)

• created_at (str)

• type (Literal['embedding'])

• usage (EmbeddingUsage | None)

• source (Literal['cache', 'api'])

Return type:

None

type: Literal['embedding']

The type of the response, must be embedding.

usage: EmbeddingUsage | None

The usage of the embedding model API invocation, if available.

source: Literal['cache', 'api']

If the response comes from the cache or the API.

class DashScopeTextEmbedding

Bases: EmbeddingModelBase

DashScope text embedding model class

__init__(api_key, model_name, embedding_cache=None)

Initialize the DashScope text embedding model class.

Parameters:

• api_key (str) – The dashscope API key.

• model_name (str) – The name of the embedding model.

• embedding_cache (EmbeddingCacheBase) – The embedding cache class instance, used to cache the embedding results to avoid repeated API calls.

Return type:

None

async __call__(text, **kwargs)

Call the DashScope embedding API.

Parameters:

• text (List[str]) – The input text to be embedded. It can be a list of strings.

• kwargs (Any)

Return type:

EmbeddingResponse

class OpenAITextEmbedding

Bases: EmbeddingModelBase

OpenAI text embedding model class.

__init__(api_key, model_name, embedding_cache=None, **kwargs)

Initialize the OpenAI text embedding model class.

Parameters:

• api_key (str) – The OpenAI API key.

• model_name (str) – The name of the embedding model.

• embedding_cache (EmbeddingCacheBase | None, defaults to None) – The embedding cache class instance, used to cache the embedding results to avoid repeated API calls.

• kwargs (Any)

Return type:

None

async __call__(text, **kwargs)

Call the OpenAI embedding API.

Parameters:

• text (List[str]) – The input text to be embedded. It can be a list of strings.

• kwargs (Any)

Return type:

EmbeddingResponse

class GeminiTextEmbedding

Bases: EmbeddingModelBase

The Gemini text embedding model.

__init__(api_key, model_name, embedding_cache=None, **kwargs)

Initialize the Gemini text embedding model class.

Parameters:

• api_key (str) – The Gemini API key.

• model_name (str) – The name of the embedding model.

• embedding_cache (EmbeddingCacheBase | None, defaults to None) – The embedding cache class instance, used to cache the embedding results to avoid repeated API calls.

• kwargs (Any)

Return type:

None

async __call__(text, **kwargs)

The Gemini embedding API call.

Parameters:

• text (List[str]) – The input text to be embedded. It can be a list of strings.

• kwargs (Any)

Return type:

EmbeddingResponse

class OllamaTextEmbedding

Bases: EmbeddingModelBase

The Ollama embedding model.

__init__(model_name, host=None, embedding_cache=None, **kwargs)

Initialize the Ollama text embedding model class.

Parameters:

• model_name (str) – The name of the embedding model.

• host (str | None, defaults to None) – The host URL for the Ollama API.

• embedding_cache (EmbeddingCacheBase | None, defaults to None) – The embedding cache class instance, used to cache the embedding results to avoid repeated API calls.

• kwargs (Any)

Return type:

None

async __call__(text, **kwargs)

Call the Ollama embedding API.

Parameters:

• text (List[str]) – The input text to be embedded. It can be a list of strings.

• kwargs (Any)

Return type:

EmbeddingResponse

class EmbeddingCacheBase

Bases: object

Base class for embedding caches, which is responsible for storing and retrieving embeddings.

abstract async store(embeddings, identifier, overwrite=False, **kwargs)

Store the embeddings with the given identifier.

Parameters:

• embeddings (List[Embedding]) – The embeddings to store.

• identifier (JSONSerializableObject) – The identifier to distinguish the embeddings.

• overwrite (bool, defaults to False) – Whether to overwrite existing embeddings with the same identifier. If True, existing embeddings will be replaced.

• kwargs (Any)

Return type:

None

abstract async retrieve(identifier)

Retrieve the embeddings with the given identifier. If not found, return None.

Parameters:

identifier (JSONSerializableObject) – The identifier to retrieve the embeddings.

Return type:

List[List[float]] | None

abstract async remove(identifier)

Remove the embeddings with the given identifier.

Parameters:

identifier (JSONSerializableObject) – The identifier to remove the embeddings.

Return type:

None

abstract async clear()

Clear all cached embeddings.

Return type:

None

class FileEmbeddingCache

Bases: EmbeddingCacheBase

The embedding cache class that stores each embeddings vector in binary files.

__init__(cache_dir='./.cache/embeddings', max_file_number=None, max_cache_size=None)

Initialize the file embedding cache class.

Parameters:

• cache_dir (str, defaults to “./.cache/embeddings”) – The directory to store the embedding files.

• max_file_number (int | None, defaults to None) – The maximum number of files to keep in the cache directory. If exceeded, the oldest files will be removed.

• max_cache_size (int | None, defaults to None) – The maximum size of the cache directory in MB. If exceeded, the oldest files will be removed until the size is within the limit.

Return type:

None

property cache_dir: str

The cache directory where the embedding files are stored.

async store(embeddings, identifier, overwrite=False, **kwargs)

Store the embeddings with the given identifier.

Parameters:

• embeddings (List[Embedding]) – The embeddings to store.

• identifier (JSONSerializableObject) – The identifier to distinguish the embeddings, which will be used to generate a hashable filename, so it should be JSON serializable (e.g. a string, number, list, dict).

• overwrite (bool, defaults to False) – Whether to overwrite existing embeddings with the same identifier. If True, existing embeddings will be replaced.

• kwargs (Any)

Return type:

None

async retrieve(identifier)

Retrieve the embeddings with the given identifier. If not found, return None.

Parameters:

identifier (JSONSerializableObject) – The identifier to retrieve the embeddings, which will be used to generate a hashable filename, so it should be JSON serializable (e.g. a string, number, list, dict).

Return type:

List[List[float]] | None

async remove(identifier)

Remove the embeddings with the given identifier.

Parameters:

identifier (JSONSerializableObject) – The identifiers to remove the embeddings, which will be used to generate a hashable filename, so it should be JSON serializable (e.g. a string, number, list, dict).

Return type:

None

async clear()

Clear the cache directory by removing all files.

Return type:

None