Skip to content

validation

mlflow_assistant.cli.validation

Validation utilities for MLflow Assistant configuration.

This module provides validation functions to check MLflow connections, AI provider configurations, and overall system setup to ensure proper operation of MLflow Assistant.

validate_mlflow_uri(uri)

Validate MLflow URI by attempting to connect.

Parameters:

Name Type Description Default
uri str

MLflow server URI

 required

Returns:

Name Type Description
bool bool

True if connection successful, False otherwise
Source code in src/mlflow_assistant/cli/validation.py 
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
def validate_mlflow_uri(uri: str) -> bool:
    """Validate MLflow URI by attempting to connect.

    Args:
        uri: MLflow server URI

    Returns:
        bool: True if connection successful, False otherwise

    """
    for endpoint in MLFLOW_VALIDATION_ENDPOINTS:
        try:
            # Try with trailing slash trimmed
            clean_uri = uri.rstrip("/")
            url = f"{clean_uri}{endpoint}"
            logger.debug(f"Trying to connect to MLflow at: {url}")

            response = requests.get(url, timeout=MLFLOW_CONNECTION_TIMEOUT)
            if response.status_code == 200:
                logger.info(f"Successfully connected to MLflow at {url}")
                return True
            logger.debug(f"Response from {url}: {response.status_code}")
        except Exception as e:
            logger.debug(f"Failed to connect to {endpoint}: {e!s}")

    # If we get here, none of the endpoints worked
    logger.warning(
        f"Could not validate MLflow at {uri} on any standard endpoint",
    )
    return False

validate_ollama_connection(uri)

Validate Ollama connection and get available models.

Parameters:

Name Type Description Default
uri str

Ollama server URI

 required

Returns:

Type Description
tuple[bool, dict[str, Any]]

Tuple[bool, Dict[str, Any]]: (is_valid, response_data)
Source code in src/mlflow_assistant/cli/validation.py 
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
def validate_ollama_connection(uri: str) -> tuple[bool, dict[str, Any]]:
    """Validate Ollama connection and get available models.

    Args:
        uri: Ollama server URI

    Returns:
        Tuple[bool, Dict[str, Any]]: (is_valid, response_data)

    """
    try:
        response = requests.get(
            f"{uri}{OLLAMA_TAGS_ENDPOINT}", timeout=OLLAMA_CONNECTION_TIMEOUT,
        )
        if response.status_code == 200:
            try:
                models_data = response.json()
                available_models = [
                    m.get("name") for m in models_data.get("models", [])
                ]
                return True, {"models": available_models}
            except Exception as e:
                logger.debug(f"Error parsing Ollama models: {e}")
                return True, {"models": []}
        else:
            return False, {}
    except Exception as e:
        logger.debug(f"Error connecting to Ollama: {e}")
        return False, {}

validate_setup(check_api_key=True)

Validate that MLflow Assistant is properly configured.

Parameters:

Name Type Description Default
check_api_key bool

Whether to check for API key if using OpenAI

 True

Returns:

Type Description
tuple[bool, str]

Tuple[bool, str]: (is_valid, error_message)
Source code in src/mlflow_assistant/cli/validation.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
def validate_setup(check_api_key: bool = True) -> tuple[bool, str]:
    """Validate that MLflow Assistant is properly configured.

    Args:
        check_api_key: Whether to check for API key if using OpenAI

    Returns:
        Tuple[bool, str]: (is_valid, error_message)

    """
    # Check MLflow URI
    mlflow_uri = get_mlflow_uri()
    if not mlflow_uri:
        return (
            False,
            "MLflow URI not configured. "
            "Run 'mlflow-assistant setup' first.",
        )

    # Get provider config
    provider_config = get_provider_config()
    if not provider_config or not provider_config.get(CONFIG_KEY_TYPE):
        return (
            False,
            "AI provider not configured. "
            "Run 'mlflow-assistant setup' first.",
        )

    # Ensure OpenAI has an API key if that's the configured provider
    if (
        check_api_key
        and provider_config.get(CONFIG_KEY_TYPE) == Provider.OPENAI.value
        and not provider_config.get(CONFIG_KEY_API_KEY)
    ):
        return (
            False,
            f"OpenAI API key not found in environment. "
            f"Set {OPENAI_API_KEY_ENV}.",
        )

    return True, ""