1  """MLflow CLI commands for Assistant integration."""
  2  
  3  import shutil
  4  import sys
  5  import threading
  6  import time
  7  from pathlib import Path
  8  
  9  import click
 10  
 11  from mlflow.assistant.config import AssistantConfig, ProjectConfig, SkillsConfig
 12  from mlflow.assistant.providers import AssistantProvider, list_providers
 13  from mlflow.assistant.providers.base import ProviderNotConfiguredError
 14  from mlflow.assistant.skill_installer import install_skills
 15  
 16  
 17  class Spinner:
 18      """Simple spinner animation for long-running operations."""
 19  
 20      def __init__(self, message: str = "Loading"):
 21          self.message = message
 22          self.spinning = False
 23          self.thread = None
 24          self.frames = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"]
 25  
 26      def _spin(self):
 27          i = 0
 28          while self.spinning:
 29              frame = self.frames[i % len(self.frames)]
 30              sys.stdout.write(f"\r{frame} {self.message}")
 31              sys.stdout.flush()
 32              time.sleep(0.1)
 33              i += 1
 34  
 35      def __enter__(self):
 36          self.spinning = True
 37          self.thread = threading.Thread(target=self._spin, name="Spinner")
 38          self.thread.start()
 39          return self
 40  
 41      def __exit__(self, *args):
 42          self.spinning = False
 43          if self.thread:
 44              self.thread.join()
 45          sys.stdout.write("\r" + " " * (len(self.message) + 4) + "\r")
 46          sys.stdout.flush()
 47  
 48  
 49  @click.command("assistant")
 50  @click.option(
 51      "--configure",
 52      is_flag=True,
 53      help="Configure or reconfigure the assistant settings",
 54  )
 55  def commands(configure: bool):
 56      """MLflow Assistant - AI-powered trace analysis.
 57  
 58      Run 'mlflow assistant --configure' to set up the assistant.
 59      """
 60      if configure:
 61          _run_configuration()
 62      else:
 63          # Check if already configured
 64          config = AssistantConfig.load()
 65          if not config.providers:
 66              click.secho(
 67                  "Assistant is not configured. Please run: mlflow assistant --configure",
 68                  fg="yellow",
 69              )
 70          else:
 71              click.secho(
 72                  "Assistant launch is not yet implemented. To use Assistant, run `mlflow assistant "
 73                  "--configure` to setup, then launch the MLflow UI manually.",
 74                  fg="yellow",
 75              )
 76  
 77  
 78  def _run_configuration():
 79      """Configure MLflow Assistant for the UI.
 80  
 81      This interactive command sets up the AI assistant feature that allows you
 82      to analyze MLflow traces directly from the UI.
 83  
 84      The command will:
 85      1. Ask which provider to use (Claude Code for now)
 86      2. Check provider availability
 87      3. Optionally connect an experiment with code repository
 88      4. Ask which model to use
 89      5. Ask where to install skills (user-level or project-level)
 90      6. Install provider-specific skills
 91      7. Save configuration
 92  
 93      Example:
 94          mlflow assistant --configure
 95      """
 96      click.echo()
 97      click.secho("╔══════════════════════════════════════════╗", fg="cyan")
 98      click.secho("║       *    .  *       .   *              ║", fg="cyan")
 99      click.secho("║   .    *  MLflow Assistant Setup   *  .  ║", fg="cyan", bold=True)
100      click.secho("║      *    .       *   .      *           ║", fg="cyan")
101      click.secho("╚══════════════════════════════════════════╝", fg="cyan")
102      click.echo()
103  
104      # Step 1: Select provider
105      provider = _prompt_provider()
106      if provider is None:
107          return
108  
109      # Step 2: Check provider availability
110      if not _check_provider(provider):
111          return
112  
113      # Step 3: Optionally connect experiment with code repository
114      project_path = _prompt_experiment_path()
115  
116      # Step 4: Ask for model
117      model = _prompt_model()
118  
119      # Step 5: Ask for skill location
120      skills_config = _prompt_skill_location(project_path)
121  
122      # Step 6: Install skills
123      skill_path = _install_skills(provider, skills_config, project_path)
124  
125      # Step 7: Save configuration
126      _save_config(provider, model, skills_config)
127  
128      # Show success message
129      _show_init_success(provider, model, skill_path)
130  
131  
132  def _prompt_provider() -> AssistantProvider | None:
133      """Prompt user to select a provider."""
134      providers = list_providers()
135  
136      click.secho("Step 1/4: Select AI Provider", fg="cyan", bold=True)
137      click.secho("-" * 30, fg="cyan")
138      click.echo()
139  
140      for i, provider in enumerate(providers, 1):
141          marker = click.style(" [recommended]", fg="green") if i == 1 else ""
142          click.echo(f"  {i}. {provider.display_name}{marker}")
143          click.secho(f"     {provider.description}", dim=True)
144  
145      click.echo()
146      click.secho("  More providers coming soon...", dim=True)
147      click.echo()
148  
149      default_provider = providers[0]
150      choice = click.prompt(
151          click.style(f"Select provider [1: {default_provider.display_name}]", fg="bright_blue"),
152          default="1",
153          type=click.Choice([str(i) for i in range(1, len(providers) + 1)]),
154          show_choices=False,
155          show_default=False,
156      )
157  
158      provider = providers[int(choice) - 1]
159      click.echo()
160      return provider
161  
162  
163  def _check_provider(provider: AssistantProvider) -> bool:
164      """Check if the selected provider is available."""
165      click.secho("Step 2/4: Checking Provider", fg="cyan", bold=True)
166      click.secho("-" * 30, fg="cyan")
167      click.echo()
168  
169      # First check if CLI is installed
170      claude_path = shutil.which("claude")
171      if not claude_path:
172          click.secho(
173              "Claude Code CLI is not installed. "
174              "Install it with: npm install -g @anthropic-ai/claude-code",
175              fg="red",
176          )
177          click.echo()
178          return False
179  
180      click.echo(f"Claude CLI found: {claude_path}")
181  
182      try:
183          spinner_msg = "Checking connection... " + click.style(
184              "(this may take a few seconds)", dim=True
185          )
186          with Spinner(spinner_msg):
187              provider.check_connection()
188          click.secho("Connection verified", fg="green")
189          click.echo()
190          return True
191      except ProviderNotConfiguredError as e:
192          click.secho(str(e), fg="red")
193          click.echo()
194          return False
195  
196  
197  def _fetch_recent_experiments(tracking_uri: str, max_results: int = 5) -> list[tuple[str, str]]:
198      """Fetch recent experiments from the tracking server.
199  
200      Returns:
201          List of (experiment_id, experiment_name) tuples.
202      """
203      import mlflow
204  
205      original_uri = mlflow.get_tracking_uri()
206      try:
207          mlflow.set_tracking_uri(tracking_uri)
208          client = mlflow.MlflowClient()
209          experiments = client.search_experiments(
210              max_results=max_results,
211              order_by=["last_update_time DESC"],
212          )
213          return [(exp.experiment_id, exp.name) for exp in experiments]
214      except Exception:
215          return []
216      finally:
217          mlflow.set_tracking_uri(original_uri)
218  
219  
220  def _resolve_experiment_id(tracking_uri: str, name_or_id: str) -> str | None:
221      """Resolve experiment name or ID to experiment ID.
222  
223      Args:
224          tracking_uri: MLflow tracking server URI.
225          name_or_id: Experiment name or ID.
226  
227      Returns:
228          Experiment ID if found, None otherwise.
229      """
230      import mlflow
231  
232      original_uri = mlflow.get_tracking_uri()
233      try:
234          mlflow.set_tracking_uri(tracking_uri)
235          client = mlflow.MlflowClient()
236  
237          # First try to get by ID (if it looks like an ID)
238          if name_or_id.isdigit():
239              try:
240                  if exp := client.get_experiment(name_or_id):
241                      return exp.experiment_id
242              except Exception:
243                  pass
244  
245          # Try to get by name
246          if exp := client.get_experiment_by_name(name_or_id):
247              return exp.experiment_id
248  
249          return None
250      except Exception:
251          return None
252      finally:
253          mlflow.set_tracking_uri(original_uri)
254  
255  
256  def _prompt_experiment_path() -> Path | None:
257      """Prompt user to optionally connect an experiment with code repository.
258  
259      Returns:
260          The project path if configured, None otherwise.
261      """
262      click.secho("Step 3/5: Experiment & Code Context ", fg="cyan", bold=True, nl=False)
263      click.secho("[Optional, Recommended]", fg="green", bold=True)
264      click.secho("-" * 30, fg="cyan")
265      click.echo()
266      click.echo("You can connect an experiment with a code repository to give")
267      click.echo("the assistant context about your source code for better analysis.")
268      click.secho("(You can also set this up later in the MLflow UI.)", dim=True)
269      click.echo()
270  
271      connect = click.confirm(
272          click.style(
273              "Do you want to connect an experiment with a code repository?", fg="bright_blue"
274          ),
275          default=True,
276      )
277  
278      if not connect:
279          click.echo()
280          return None
281  
282      click.echo()
283  
284      # Ask for tracking URI to fetch experiments
285      tracking_uri = click.prompt(
286          click.style("Enter the MLflow tracking server URI", fg="bright_blue"),
287          default="http://localhost:5000",
288      )
289  
290      click.echo()
291      click.secho("Fetching recent experiments...", dim=True)
292  
293      # Fetch recent experiments
294      experiments = _fetch_recent_experiments(tracking_uri)
295  
296      if not experiments:
297          click.secho("Could not fetch experiments from the server.", fg="yellow")
298          click.echo("You can set this up later in the MLflow UI.")
299          click.echo()
300          return None
301  
302      click.echo()
303      click.echo(click.style("Select an experiment to connect:", fg="bright_blue"))
304      click.echo()
305  
306      for i, (exp_id, exp_name) in enumerate(experiments, 1):
307          click.echo(f"  {i}. {exp_name} (ID: {exp_id})")
308  
309      other_option = len(experiments) + 1
310      click.echo(f"  {other_option}. Enter experiment name or ID manually")
311      click.echo()
312  
313      choice = click.prompt(
314          click.style("Select experiment", fg="bright_blue"),
315          type=click.IntRange(1, other_option),
316          default=1,
317      )
318  
319      if choice == other_option:
320          while True:
321              click.echo()
322              name_or_id = click.prompt(
323                  click.style("Experiment name or ID", fg="bright_blue"), default=""
324              )
325              if not name_or_id:
326                  click.secho("No experiment specified. Please try again.", fg="yellow")
327                  continue
328  
329              experiment_id = _resolve_experiment_id(tracking_uri, name_or_id)
330              if experiment_id:
331                  # Use the input as display name (could be name or ID)
332                  experiment_name = name_or_id
333                  break
334  
335              click.secho(
336                  f"Experiment '{name_or_id}' not found. Please try again.",
337                  fg="red",
338              )
339      else:
340          experiment_id, experiment_name = experiments[choice - 1]
341  
342      click.secho(
343          f"Experiment '{experiment_name}' selected",
344          fg="green",
345      )
346      click.echo()
347  
348      # Ask for project path
349      default_path = str(Path.cwd())
350      while True:
351          raw_path = click.prompt(
352              click.style("Enter the path to your project directory:", fg="bright_blue"),
353              default=default_path,
354          )
355          # Expand ~ and resolve relative paths
356          expanded_path = Path(raw_path).expanduser().resolve()
357          if expanded_path.is_dir():
358              project_path = str(expanded_path)
359              break
360          click.secho(f"Directory '{raw_path}' does not exist. Please try again.", fg="red")
361  
362      # Save the project path mapping locally
363      try:
364          config = AssistantConfig.load()
365          config.projects[experiment_id] = ProjectConfig(type="local", location=project_path)
366          config.save()
367          click.secho(
368              f"Project path {project_path} is saved for experiment '{experiment_name}'",
369              fg="green",
370          )
371      except Exception as e:
372          click.secho(f"Error saving project path: {e}", fg="red")
373  
374      click.echo()
375      return expanded_path
376  
377  
378  def _prompt_model() -> str:
379      """Prompt user for model selection."""
380      click.secho("Step 4/5: Model Selection", fg="cyan", bold=True)
381      click.secho("-" * 30, fg="cyan")
382      click.echo()
383      click.echo("Choose a model for analysis:")
384      click.secho("  - Press Enter to use the default model (recommended)", dim=True)
385      click.secho("  - Or type a specific model name (e.g., claude-sonnet-4-20250514)", dim=True)
386      click.echo()
387  
388      model = click.prompt(click.style("Model", fg="bright_blue"), default="default")
389      click.echo()
390      return model
391  
392  
393  def _prompt_skill_location(project_path: Path | None) -> SkillsConfig:
394      """Prompt user for skill installation location.
395  
396      Args:
397          project_path: The project path from experiment setup, or None if skipped.
398  
399      Returns:
400          SkillsConfig with the selected location type and optional custom path.
401      """
402      click.secho("Step 5/5: Skill Installation Location", fg="cyan", bold=True)
403      click.secho("-" * 30, fg="cyan")
404      click.echo()
405      click.echo("Choose where to install MLflow skills for Assistant:")
406      click.echo()
407  
408      # TODO: Update this when we support other providers
409      user_path = Path.home() / ".claude" / "skills"
410      click.echo(f"  1. User level ({user_path})")
411      click.secho("     Skills available globally across all projects", dim=True)
412      click.echo()
413  
414      if project_path:
415          project_skill_path = project_path / ".claude" / "skills"
416          click.echo(f"  2. Project level ({project_skill_path})")
417          click.secho("     Skills available only in this project", dim=True)
418          click.echo()
419          click.echo("  3. Custom location")
420          click.secho("     Specify a custom path for skills", dim=True)
421          click.echo()
422          valid_choices = ["1", "2", "3"]
423      else:
424          click.echo("  2. Custom location")
425          click.secho("     Specify a custom path for skills", dim=True)
426          click.echo()
427          valid_choices = ["1", "2"]
428  
429      choice = click.prompt(
430          click.style("Select location [1: User level]", fg="bright_blue"),
431          default="1",
432          type=click.Choice(valid_choices),
433          show_choices=False,
434          show_default=False,
435      )
436  
437      click.echo()
438  
439      if choice == "1":
440          return SkillsConfig(type="global")
441      elif choice == "2" and project_path:
442          return SkillsConfig(type="project")
443      else:
444          # Custom location
445          while True:
446              raw_path = click.prompt(
447                  click.style("Enter the custom path for skills", fg="bright_blue"),
448                  default=str(user_path),
449              )
450              expanded_path = Path(raw_path).expanduser().resolve()
451              # For custom paths, we'll create the directory, so just check parent exists
452              if expanded_path.parent.exists() or expanded_path.exists():
453                  click.echo()
454                  return SkillsConfig(type="custom", custom_path=str(expanded_path))
455              click.secho(
456                  f"Parent directory '{expanded_path.parent}' does not exist. Please try again.",
457                  fg="red",
458              )
459  
460  
461  def _install_skills(
462      provider: AssistantProvider, skills_config: SkillsConfig, project_path: Path | None
463  ) -> Path:
464      """Install skills bundled with MLflow.
465  
466      Returns:
467          The resolved path where skills were installed.
468      """
469      match skills_config.type:
470          case "global":
471              skill_path = provider.resolve_skills_path(Path.home())
472          case "project":
473              skill_path = provider.resolve_skills_path(project_path)
474          case "custom":
475              skill_path = Path(skills_config.custom_path).expanduser()
476      if installed_skills := install_skills(skill_path):
477          for skill in installed_skills:
478              click.secho(f"  - {skill}")
479      else:
480          click.secho("No skills available to install.", fg="yellow")
481      click.echo()
482      return skill_path
483  
484  
485  def _save_config(provider: AssistantProvider, model: str, skills_config: SkillsConfig) -> None:
486      """Save configuration to file."""
487      click.secho("Saving Configuration", fg="cyan", bold=True)
488      click.secho("-" * 30, fg="cyan")
489  
490      config = AssistantConfig.load()
491      config.set_provider(provider.name, model)
492      config.providers[provider.name].skills = skills_config
493      config.save()
494  
495      click.secho("Configuration saved", fg="green")
496      click.echo()
497  
498  
499  def _show_init_success(provider: AssistantProvider, model: str, skill_path: Path) -> None:
500      """Show success message and next steps."""
501      click.secho("  ~ * ~ * ~ * ~ * ~ * ~ * ~ * ~", fg="green")
502      click.secho("        Setup Complete!        ", fg="green", bold=True)
503      click.secho("  ~ * ~ * ~ * ~ * ~ * ~ * ~ * ~", fg="green")
504      click.echo()
505      click.secho("Configuration:", bold=True)
506      click.echo(f"  Provider: {provider.display_name}")
507      click.echo(f"  Model: {model}")
508      click.echo(f"  Skills: {skill_path}")
509      click.echo()
510      click.secho("Next steps:", bold=True)
511      click.echo("  1. Start MLflow server:")
512      click.secho("     $ mlflow server", fg="cyan")
513      click.echo()
514      click.echo("  2. Open MLflow UI and navigate to an experiment")
515      click.echo()
516      click.echo("  3. Click 'Ask Assistant'")
517      click.echo()
518      click.secho("To reconfigure, run: ", nl=False)
519      click.secho("mlflow assistant --configure", fg="cyan")