Skip to content

core

mlflow_assistant.core

Core functionality for MLflow Assistant.

This subpackage contains the core modules for managing connections, workflows, and interactions with the MLflow Tracking Server.

cli

Command-line interface (CLI) for MLflow Assistant.

This module provides the CLI entry points for interacting with the MLflow Assistant, allowing users to manage connections, workflows, and other operations via the command line.

connection

MLflow connection module for handling connections to MLflow Tracking Server.

This module provides functionality to connect to both local and remote MLflow Tracking Servers using environment variables or direct configuration.

MLflowConnection(tracking_uri=None, client_factory=None)

MLflow connection class to handle connections to MLflow Tracking Server.

This class provides functionality to connect to both local and remote MLflow Tracking Servers.

Initialize MLflow connection.

Parameters:

Name Type Description Default
tracking_uri str | None

URI of the MLflow Tracking Server. If None, will try to get from environment.

 None
client_factory Any

A callable to create the MlflowClient instance. Defaults to MlflowClient.

 None
Source code in src/mlflow_assistant/core/connection.py 
31
32
33
34
35
36
37
38
39
40
41
42
def __init__(self, tracking_uri: str | None = None, client_factory: Any = None):
    """Initialize MLflow connection.

    Args:
        tracking_uri: URI of the MLflow Tracking Server. If None, will try to get from environment.
        client_factory: A callable to create the MlflowClient instance. Defaults to MlflowClient.

    """
    self.config = self._load_config(tracking_uri=tracking_uri)
    self.client = None
    self.is_connected_flag = False
    self.client_factory = client_factory or MlflowClient
connect()

Connect to MLflow Tracking Server.

Returns
Tuple[bool, str]: (success, message)
Source code in src/mlflow_assistant/core/connection.py 
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
def connect(self) -> tuple[bool, str]:
    """Connect to MLflow Tracking Server.

    Returns
    -------
        Tuple[bool, str]: (success, message)

    """
    try:
        logger.debug(f"Connecting to MLflow Tracking Server at {self.config.tracking_uri}")
        mlflow.set_tracking_uri(self.config.tracking_uri)
        self.client = self.client_factory(tracking_uri=self.config.tracking_uri)
        self.client.search_experiments()  # Trigger connection attempt
        self.is_connected_flag = True
        logger.debug(f"Successfully connected to MLflow Tracking Server at {self.config.tracking_uri}")
        return True, f"Successfully connected to MLflow Tracking Server at {self.config.tracking_uri}"
    except Exception as e:
        self.is_connected_flag = False
        logger.exception(f"Failed to connect to MLflow Tracking Server: {e}")
        return False, f"Failed to connect to MLflow Tracking Server: {e!s}"
get_client()

Get MLflow client instance.

Returns
MlflowClient: MLflow client instance.
Raises
MLflowConnectionError: If not connected to MLflow Tracking Server.
Source code in src/mlflow_assistant/core/connection.py 
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
def get_client(self) -> MlflowClient:
    """Get MLflow client instance.

    Returns
    -------
        MlflowClient: MLflow client instance.

    Raises
    ------
        MLflowConnectionError: If not connected to MLflow Tracking Server.

    """
    if self.client is None:
        msg = "Not connected to MLflow Tracking Server. Call connect() first."
        raise MLflowConnectionError(msg)
    return self.client
get_connection_info()

Get connection information.

Returns
Dict[str, Any]: Connection information.
Source code in src/mlflow_assistant/core/connection.py 
105
106
107
108
109
110
111
112
113
114
115
116
117
def get_connection_info(self) -> dict[str, Any]:
    """Get connection information.

    Returns
    -------
        Dict[str, Any]: Connection information.

    """
    return {
        "tracking_uri": self.config.tracking_uri,
        "connection_type": self.config.connection_type,
        "is_connected": self.is_connected_flag,
    }
is_connected()

Check if connected to MLflow Tracking Server.

Returns
bool: True if connected, False otherwise.
Source code in src/mlflow_assistant/core/connection.py 
 95
 96
 97
 98
 99
100
101
102
103
def is_connected(self) -> bool:
    """Check if connected to MLflow Tracking Server.

    Returns
    -------
        bool: True if connected, False otherwise.

    """
    return self.is_connected_flag

core

Core utilities and functionality for MLflow Assistant.

This module provides foundational classes, functions, and utilities used across the MLflow Assistant project, including shared logic for managing workflows and interactions with the MLflow Tracking Server.

get_mlflow_client()

Initialize and return an MLflow client instance.

Returns
MlflowClient: An instance of the MLflow client.
Source code in src/mlflow_assistant/core/core.py 
10
11
12
13
14
15
16
17
18
def get_mlflow_client():
    """Initialize and return an MLflow client instance.

    Returns
    -------
        MlflowClient: An instance of the MLflow client.

    """
    return MlflowClient()

provider

Provider integrations for MLflow Assistant.

This module defines the interfaces and implementations for integrating with various large language model (LLM) providers, such as OpenAI and Ollama.

workflow

Workflow management for LangGraph in MLflow Assistant.

This module provides functionality for defining, managing, and executing workflows using LangGraph, enabling seamless integration with MLflow for tracking and managing machine learning workflows.