Skip to content

commands

mlflow_assistant.cli.commands

CLI commands for MLflow Assistant.

This module contains the main CLI commands for interacting with MLflow using natural language queries through various AI providers.

cli(verbose)

MLflow Assistant: Interact with MLflow using LLMs.

This CLI tool helps you to interact with MLflow using natural language.

Source code in src/mlflow_assistant/cli/commands.py
@click.group()
@click.option("--verbose", "-v", is_flag=True, help="Enable verbose logging")
def cli(verbose):
    """MLflow Assistant: Interact with MLflow using LLMs.

    This CLI tool helps you to interact with MLflow using natural language.
    """
    # Configure logging
    log_level = logging.DEBUG if verbose else logging.INFO
    logging.basicConfig(level=log_level, format=LOG_FORMAT)

mock_process_query(query, provider_config, verbose=False)

Mock function that simulates the query processing workflow.

This will be replaced with the actual implementation later.

Parameters:

Name Type Description Default
query str

The user's query

required
provider_config dict[str, Any]

The AI provider configuration

required
verbose bool

Whether to show verbose output

False

Returns:

Type Description
dict[str, Any]

Dictionary with mock response information

Source code in src/mlflow_assistant/cli/commands.py
def mock_process_query(
    query: str, provider_config: dict[str, Any], verbose: bool = False,
) -> dict[str, Any]:
    """Mock function that simulates the query processing workflow.

    This will be replaced with the actual implementation later.

    Args:
        query: The user's query
        provider_config: The AI provider configuration
        verbose: Whether to show verbose output

    Returns:
        Dictionary with mock response information

    """
    # Create a mock response
    provider_type = provider_config.get(
        CONFIG_KEY_TYPE, DEFAULT_STATUS_NOT_CONFIGURED,
    )
    model = provider_config.get(
        CONFIG_KEY_MODEL, DEFAULT_STATUS_NOT_CONFIGURED,
    )

    response_text = (
        f"This is a mock response to: '{query}'\n\n"
        f"The MLflow integration will be implemented soon!"
    )

    if verbose:
        response_text += f"\n\nDebug: Using {provider_type} with {model}"

    return {
        "original_query": query,
        "provider_config": {
            CONFIG_KEY_TYPE: provider_type,
            CONFIG_KEY_MODEL: model,
        },
        "enhanced": False,
        "response": response_text,
    }

setup()

Run the interactive setup wizard.

This wizard helps you configure MLflow Assistant.

Source code in src/mlflow_assistant/cli/commands.py
@cli.command()
def setup():
    """Run the interactive setup wizard.

    This wizard helps you configure MLflow Assistant.
    """
    setup_wizard()

start(verbose)

Start an interactive chat session with MLflow Assistant.

This opens an interactive chat session where you can ask questions about your MLflow experiments, models, and data. Type /bye to exit the session.

Examples of questions you can ask: - What are my best performing models for classification? - Show me details of experiment 'customer_churn' - Compare runs abc123 and def456 - Which hyperparameters should I try next for my regression model?

Commands: - /bye: Exit the chat session - /help: Show help about available commands - /clear: Clear the screen

Source code in src/mlflow_assistant/cli/commands.py
@cli.command()
@click.option("--verbose", "-v", is_flag=True, help="Show verbose output")
def start(verbose):
    """Start an interactive chat session with MLflow Assistant.

    This opens an interactive chat session where you can ask questions about
    your MLflow experiments, models, and data. Type /bye to exit the session.

    Examples of questions you can ask:
    - What are my best performing models for classification?
    - Show me details of experiment 'customer_churn'
    - Compare runs abc123 and def456
    - Which hyperparameters should I try next for my regression model?

    Commands:
    - /bye: Exit the chat session
    - /help: Show help about available commands
    - /clear: Clear the screen
    """
    # Use validation function to check setup
    is_valid, error_message = validate_setup()
    if not is_valid:
        click.echo(f"❌ Error: {error_message}")
        return

    # Get provider config
    provider_config = get_provider_config()

    # Print welcome message and instructions
    provider_type = provider_config.get(
        CONFIG_KEY_TYPE, DEFAULT_STATUS_NOT_CONFIGURED,
        )
    model = provider_config.get(
        CONFIG_KEY_MODEL, DEFAULT_STATUS_NOT_CONFIGURED,
        )

    click.echo("\n🤖 MLflow Assistant Chat Session")
    click.echo(f"Connected to MLflow at: {get_mlflow_uri()}")
    click.echo(f"Using {provider_type.upper()} with model: {model}")
    click.echo("\nType your questions and press Enter.")
    click.echo(f"Type {Command.EXIT.value} to exit.")
    click.echo("=" * 70)

    # Start interactive loop
    while True:
        # Get user input with a prompt
        try:
            query = click.prompt("\n🧑", prompt_suffix="").strip()
        except (KeyboardInterrupt, EOFError):
            click.echo("\nExiting chat session...")
            break

        # Handle special commands
        action = _handle_special_commands(query)
        if action == "exit":
            break
        if action == "continue":
            continue

        # Process the query
        asyncio.run(_process_user_query(query, provider_config, verbose))

version()

Show MLflow Assistant version information.

Source code in src/mlflow_assistant/cli/commands.py
@cli.command()
def version():
    """Show MLflow Assistant version information."""
    from mlflow_assistant import __version__

    click.echo(f"MLflow Assistant version: {__version__}")

    # Show configuration
    config = load_config()
    mlflow_uri = config.get(
        CONFIG_KEY_MLFLOW_URI, DEFAULT_STATUS_NOT_CONFIGURED,
        )
    provider = config.get(CONFIG_KEY_PROVIDER, {}).get(
        CONFIG_KEY_TYPE, DEFAULT_STATUS_NOT_CONFIGURED,
    )
    model = config.get(CONFIG_KEY_PROVIDER, {}).get(
        CONFIG_KEY_MODEL, DEFAULT_STATUS_NOT_CONFIGURED,
    )

    click.echo(f"MLflow URI: {mlflow_uri}")
    click.echo(f"Provider: {provider}")
    click.echo(f"Model: {model}")