dataflow.registries.model_registry

Model Registry Client for Haive.

This module provides a client interface for working with the registry system to access LLM and embedding models with dynamic environment variable detection.

Classes

ModelRegistry

Client for working with registered LLM and embedding models.

Module Contents

class dataflow.registries.model_registry.ModelRegistry

Client for working with registered LLM and embedding models.

Initialize the model registry client.

detect_environment_variables()

Detect available environment variables for LLM and embedding. providers.

Returns:

Dict mapping provider names to available environment variables

get_available_embedding_providers()

Get all available embedding providers.

Returns:

List of available embedding providers

Return type:

list[dict[str, Any]]

get_available_llm_providers()

Get all available LLM providers.

Returns:

List of available LLM providers

Return type:

list[dict[str, Any]]

get_embedding_model(model_id)

Get a specific embedding model by ID.

Parameters:

model_id (str) – ID of the model

Returns:

Model data or None if not found

Return type:

dict[str, Any] | None

get_embedding_models(provider=None, only_available=False)

Get all embedding models, optionally filtered by provider.

Parameters:

• provider (str | None) – Optional provider name to filter by

• only_available (bool) – If True, only return models from available providers

Returns:

List of embedding model data

Return type:

list[dict[str, Any]]

get_llm_model(model_id)

Get a specific LLM model by ID.

Parameters:

model_id (str) – ID of the model

Returns:

Model data or None if not found

Return type:

dict[str, Any] | None

get_llm_models(provider=None, only_available=False)

Get all LLM models, optionally filtered by provider.

Parameters:

• provider (str | None) – Optional provider name to filter by

• only_available (bool) – If True, only return models from available providers

Returns:

List of LLM model data

Return type:

list[dict[str, Any]]

get_required_environment_vars()

Scan source code to detect environment variables used by LLM. providers.

Returns:

List of environment variables with provider mapping

Return type:

list[dict[str, Any]]

get_secret_from_vault(var_name)

Get a secret value from the vault.

Parameters:

var_name (str) – Name of the environment variable

Returns:

Secret value or None if not found

Return type:

str | None

normalize_model_data(model)

Normalize model data to ensure consistent field access.

Parameters:

model (dict[str, Any]) – Model data to normalize

Returns:

Normalized model data

Return type:

dict[str, Any]

update_provider_availability()

Scan environment variables and update provider availability. status.