Class AlloyDBModelManager (0.4.0)

AlloyDBModelManager(
    key: object, engine: llama_index_alloydb_pg.engine.AlloyDBEngine
)

Register, inspect and drop AlloyDB AI model endpoints.

See https://cloud.google.com/alloydb/docs/ai/register-model-endpoint.

Methods

AlloyDBModelManager

AlloyDBModelManager(
    key: object, engine: llama_index_alloydb_pg.engine.AlloyDBEngine
)

AlloyDBModelManager constructor.

Parameters
Name Description
key object

Prevent direct constructor usage.

engine AlloyDBEngine

AlloyDB engine with a connection pool.

Exceptions
Type Description
Exception If the constructor is called directly by the user.

__acreate_model

__acreate_model(**args: typing.Optional[str]) -> None

Call google_ml.create_model with the arguments that are not None.

__adrop_model

__adrop_model(model_id: str) -> None

Drop a model endpoint with google_ml.drop_model.

__aget_model

__aget_model(
    model_id: str,
) -> typing.Optional[llama_index_alloydb_pg.model_manager.AlloyDBModel]

Get a registered model endpoint.

__alist_models

__alist_models() -> list[llama_index_alloydb_pg.model_manager.AlloyDBModel]

List all model endpoints.

__avalidate

__avalidate() -> None

Raise ValueError if model endpoint management is unavailable.

acreate_model

acreate_model(
    model_id: str,
    *,
    model_provider: str,
    model_type: str,
    model_qualified_name: typing.Optional[str] = None,
    model_request_url: typing.Optional[str] = None,
    model_auth_type: typing.Optional[str] = None,
    model_auth_id: typing.Optional[str] = None,
    generate_headers_fn: typing.Optional[str] = None,
    model_in_transform_fn: typing.Optional[str] = None,
    model_out_transform_fn: typing.Optional[str] = None
) -> None

Register a model endpoint. Arguments left as None use server defaults.

See https://cloud.google.com/alloydb/docs/reference/model-endpoint-reference for the values each provider and type accepts.

Parameters
Name Description
model_id str

A unique ID for the model endpoint.

model_provider str

The provider, for example "google" or "custom".

model_type str

The type, for example "text_embedding" or "generic".

model_qualified_name Optional[str]

The qualified model name. With provider "google" and an embedding type, it sets the default request URL. Leaving it and model_request_url unset registers an unusable endpoint.

model_request_url Optional[str]

The model-specific HTTPS endpoint URL. With provider "google", types other than "text_embedding" and "multimodal_embedding" (for example "generic") need it.

model_auth_type Optional[str]

"alloydb_service_agent_iam" (for Vertex AI models, model_auth_id must be None) or "secret_manager" (model_auth_id is required).

model_auth_id Optional[str]

The secret ID registered with google_ml.create_sm_secret, for "secret_manager".

generate_headers_fn Optional[str]

SQL function for custom headers.

model_in_transform_fn Optional[str]

SQL function for the input.

model_out_transform_fn Optional[str]

SQL function for the output.

Exceptions
Type Description
ValueError If model_id, model_provider or model_type is not a non-blank string, or another argument is given but is not one.
DBAPIError If the database rejects the call (e.g. duplicate model_id).

adrop_model

adrop_model(model_id: str) -> None

Drop a registered model endpoint.

Parameter
Name Description
model_id str

The model ID. Unknown or built-in IDs are a silent no-op.

aget_model

aget_model(
    model_id: str,
) -> typing.Optional[llama_index_alloydb_pg.model_manager.AlloyDBModel]

Get a registered model endpoint.

Parameter
Name Description
model_id str

The ID of the model endpoint.

Returns
Type Description
Optional[AlloyDBModel] The model endpoint, or None if it is missing.

alist_models

alist_models() -> list[llama_index_alloydb_pg.model_manager.AlloyDBModel]

List all model endpoints, including the built-in ones.

Returns
Type Description
list[AlloyDBModel] The model endpoints ordered by model ID.

create

create(
    engine: llama_index_alloydb_pg.engine.AlloyDBEngine,
) -> llama_index_alloydb_pg.model_manager.AlloyDBModelManager

Create an AlloyDBModelManager after checking model endpoint support.

Parameter
Name Description
engine AlloyDBEngine

AlloyDB engine with a connection pool.

Exceptions
Type Description
ValueError If model endpoint management is not available.
Returns
Type Description
AlloyDBModelManager A new model manager.

create_model

create_model(
    model_id: str,
    *,
    model_provider: str,
    model_type: str,
    model_qualified_name: typing.Optional[str] = None,
    model_request_url: typing.Optional[str] = None,
    model_auth_type: typing.Optional[str] = None,
    model_auth_id: typing.Optional[str] = None,
    generate_headers_fn: typing.Optional[str] = None,
    model_in_transform_fn: typing.Optional[str] = None,
    model_out_transform_fn: typing.Optional[str] = None
) -> None

Synchronously register a model endpoint.

Takes the same arguments and raises the same errors as acreate_model.

create_sync

create_sync(
    engine: llama_index_alloydb_pg.engine.AlloyDBEngine,
) -> llama_index_alloydb_pg.model_manager.AlloyDBModelManager

Synchronously create an AlloyDBModelManager after checking support.

Parameter
Name Description
engine AlloyDBEngine

AlloyDB engine with a connection pool.

Exceptions
Type Description
ValueError If model endpoint management is not available.
Returns
Type Description
AlloyDBModelManager A new model manager.

drop_model

drop_model(model_id: str) -> None

Synchronously drop a registered model endpoint.

Parameter
Name Description
model_id str

The model ID. Unknown or built-in IDs are a silent no-op.

get_model

get_model(
    model_id: str,
) -> typing.Optional[llama_index_alloydb_pg.model_manager.AlloyDBModel]

Synchronously get a registered model endpoint.

Parameter
Name Description
model_id str

The ID of the model endpoint.

Returns
Type Description
Optional[AlloyDBModel] The model endpoint, or None if it is missing.

list_models

list_models() -> list[llama_index_alloydb_pg.model_manager.AlloyDBModel]

Synchronously list all model endpoints, including the built-in ones.

Returns
Type Description
list[AlloyDBModel] The model endpoints ordered by model ID.