Skip to content

config

mlflow_assistant.utils.config

Configuration management utilities for MLflow Assistant.

This module provides functions for loading, saving, and accessing configuration settings for MLflow Assistant, including MLflow URI and AI provider settings. Configuration is stored in YAML format in the user's home directory.

ensure_config_dir()

Ensure the configuration directory exists.

Source code in src/mlflow_assistant/utils/config.py 
43
44
45
46
47
def ensure_config_dir():
    """Ensure the configuration directory exists."""
    if not CONFIG_DIR.exists():
        CONFIG_DIR.mkdir(parents=True)
        logger.info(f"Created configuration directory at {CONFIG_DIR}")

get_mlflow_uri()

Get the MLflow URI from config or environment.

Returns:

Type Description
str | None

Optional[str]: The MLflow URI or None if not configured
Source code in src/mlflow_assistant/utils/config.py 
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
def get_mlflow_uri() -> str | None:
    """Get the MLflow URI from config or environment.

    Returns:
        Optional[str]: The MLflow URI or None if not configured

    """
    # Environment variable should take precedence
    if mlflow_uri_env := os.environ.get(MLFLOW_URI_ENV):
        return mlflow_uri_env

    # Fall back to config
    config = load_config()
    return config.get(CONFIG_KEY_MLFLOW_URI)

get_provider_config()

Get the AI provider configuration.

Returns:

Type Description
dict[str, Any]

Dict[str, Any]: The provider configuration
Source code in src/mlflow_assistant/utils/config.py 
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
def get_provider_config() -> dict[str, Any]:
    """Get the AI provider configuration.

    Returns:
        Dict[str, Any]: The provider configuration

    """
    config = load_config()
    provider = config.get(CONFIG_KEY_PROVIDER, {})

    provider_type = provider.get(CONFIG_KEY_TYPE)

    if provider_type == Provider.OPENAI.value:
        # Environment variable should take precedence
        api_key = (os.environ.get(OPENAI_API_KEY_ENV) or
                   provider.get(CONFIG_KEY_API_KEY))
        return {
            CONFIG_KEY_TYPE: Provider.OPENAI.value,
            CONFIG_KEY_API_KEY: api_key,
            CONFIG_KEY_MODEL: provider.get(
                CONFIG_KEY_MODEL, Provider.get_default_model(Provider.OPENAI),
            ),
        }

    if provider_type == Provider.OLLAMA.value:
        return {
            CONFIG_KEY_TYPE: Provider.OLLAMA.value,
            CONFIG_KEY_URI: provider.get(CONFIG_KEY_URI, DEFAULT_OLLAMA_URI),
            CONFIG_KEY_MODEL: provider.get(
                CONFIG_KEY_MODEL, Provider.get_default_model(Provider.OLLAMA),
            ),
        }

    if provider_type == Provider.DATABRICKS.value:
        # Set environment variables for Databricks profile
        _set_environment_variables(provider.get(CONFIG_KEY_PROFILE))

        return {
            CONFIG_KEY_TYPE: Provider.DATABRICKS.value,
            CONFIG_KEY_PROFILE: provider.get(CONFIG_KEY_PROFILE),
            CONFIG_KEY_MODEL: provider.get(CONFIG_KEY_MODEL),
        }

    return {CONFIG_KEY_TYPE: None}

load_config()

Load configuration from file.

Source code in src/mlflow_assistant/utils/config.py 
50
51
52
53
54
55
56
57
58
59
60
61
62
63
def load_config() -> dict[str, Any]:
    """Load configuration from file."""
    if not CONFIG_FILE.exists():
        logger.info(f"No configuration file found at {CONFIG_FILE}")
        return {}

    try:
        with open(CONFIG_FILE) as f:
            config = yaml.safe_load(f) or {}
            logger.debug(f"Loaded configuration: {config}")
            return config
    except Exception as e:
        logger.error(f"Error loading configuration: {e}")
        return {}

save_config(config)

Save configuration to file.

Parameters:

Name Type Description Default
config dict[str, Any]

Configuration dictionary to save

 required

Returns:

Name Type Description
bool bool

True if successful, False otherwise
Source code in src/mlflow_assistant/utils/config.py 
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
def save_config(config: dict[str, Any]) -> bool:
    """Save configuration to file.

    Args:
        config: Configuration dictionary to save

    Returns:
        bool: True if successful, False otherwise

    """
    ensure_config_dir()

    try:
        with open(CONFIG_FILE, "w") as f:
            yaml.dump(config, f)
        logger.info(f"Configuration saved to {CONFIG_FILE}")
        return True
    except Exception as e:
        logger.error(f"Error saving configuration: {e}")
        return False