1  """MLflow CLI commands for Claude Code integration."""
  2  
  3  from pathlib import Path
  4  
  5  import click
  6  
  7  from mlflow.claude_code.config import get_tracing_status, setup_environment_config
  8  from mlflow.claude_code.hooks import disable_tracing_hooks, setup_hooks_config, stop_hook_handler
  9  
 10  
 11  @click.group("autolog")
 12  def commands():
 13      """Commands for autologging with MLflow."""
 14  
 15  
 16  @commands.group("claude", invoke_without_command=True)
 17  @click.option(
 18      "--directory",
 19      "-d",
 20      default=".",
 21      type=click.Path(file_okay=False, dir_okay=True),
 22      help="Directory to set up tracing in (default: current directory)",
 23  )
 24  @click.option(
 25      "--tracking-uri", "-u", help="MLflow tracking URI (e.g., 'databricks' or 'file://mlruns')"
 26  )
 27  @click.option("--experiment-id", "-e", help="MLflow experiment ID")
 28  @click.option("--experiment-name", "-n", help="MLflow experiment name")
 29  @click.option("--disable", is_flag=True, help="Disable Claude tracing in the specified directory")
 30  @click.option("--status", is_flag=True, help="Show current tracing status")
 31  @click.pass_context
 32  def claude(
 33      ctx: click.Context,
 34      directory: str,
 35      tracking_uri: str | None,
 36      experiment_id: str | None,
 37      experiment_name: str | None,
 38      disable: bool,
 39      status: bool,
 40  ) -> None:
 41      """Set up Claude Code tracing in a directory.
 42  
 43      This command configures Claude Code hooks to automatically trace conversations
 44      to MLflow. After setup, use the regular 'claude' command and traces will be
 45      automatically created.
 46  
 47      Examples:
 48  
 49        # Set up tracing in current directory with local storage
 50        mlflow autolog claude
 51  
 52        # Set up tracing in a specific project directory
 53        mlflow autolog claude -d ~/my-project
 54  
 55        # Set up tracing with Databricks
 56        mlflow autolog claude -u databricks -e 123456789
 57  
 58        # Set up tracing with custom tracking URI
 59        mlflow autolog claude -u file://./custom-mlruns
 60  
 61        # Disable tracing in current directory
 62        mlflow autolog claude --disable
 63      """
 64      # Skip setup when a subcommand (e.g., stop-hook) is being invoked
 65      if ctx.invoked_subcommand is not None:
 66          return
 67  
 68      target_dir = Path(directory).resolve()
 69      claude_dir = target_dir / ".claude"
 70      settings_file = claude_dir / "settings.json"
 71  
 72      if status:
 73          _show_status(target_dir, settings_file)
 74          return
 75  
 76      if disable:
 77          _handle_disable(settings_file)
 78          return
 79  
 80      click.echo(f"Configuring Claude tracing in: {target_dir}")
 81  
 82      # Create .claude directory and set up hooks
 83      claude_dir.mkdir(parents=True, exist_ok=True)
 84      setup_hooks_config(settings_file)
 85      click.echo("✅ Claude Code hooks configured")
 86  
 87      # Set up environment variables
 88      setup_environment_config(settings_file, tracking_uri, experiment_id, experiment_name)
 89  
 90      # Show final status
 91      _show_setup_status(target_dir, tracking_uri, experiment_id, experiment_name)
 92  
 93  
 94  def _handle_disable(settings_file: Path) -> None:
 95      """Handle disable command."""
 96      if disable_tracing_hooks(settings_file):
 97          click.echo("✅ Claude tracing disabled")
 98      else:
 99          click.echo("❌ No Claude configuration found - tracing was not enabled")
100  
101  
102  def _show_status(target_dir: Path, settings_file: Path) -> None:
103      """Show current tracing status."""
104      click.echo(f"📍 Claude tracing status in: {target_dir}")
105  
106      status = get_tracing_status(settings_file)
107  
108      if not status.enabled:
109          click.echo("❌ Claude tracing is not enabled")
110          if status.reason:
111              click.echo(f"   Reason: {status.reason}")
112          return
113  
114      click.echo("✅ Claude tracing is ENABLED")
115      click.echo(f"📊 Tracking URI: {status.tracking_uri}")
116  
117      if status.experiment_id:
118          click.echo(f"🔬 Experiment ID: {status.experiment_id}")
119      elif status.experiment_name:
120          click.echo(f"🔬 Experiment Name: {status.experiment_name}")
121      else:
122          click.echo("🔬 Experiment: Default (experiment 0)")
123  
124  
125  def _show_setup_status(
126      target_dir: Path,
127      tracking_uri: str | None,
128      experiment_id: str | None,
129      experiment_name: str | None,
130  ) -> None:
131      """Show setup completion status."""
132      current_dir = Path.cwd().resolve()
133  
134      click.echo("\n" + "=" * 50)
135      click.echo("🎯 Claude Tracing Setup Complete!")
136      click.echo("=" * 50)
137  
138      click.echo(f"📁 Directory: {target_dir}")
139  
140      # Show tracking configuration
141      if tracking_uri:
142          click.echo(f"📊 Tracking URI: {tracking_uri}")
143  
144      if experiment_id:
145          click.echo(f"🔬 Experiment ID: {experiment_id}")
146      elif experiment_name:
147          click.echo(f"🔬 Experiment Name: {experiment_name}")
148      else:
149          click.echo("🔬 Experiment: Default (experiment 0)")
150  
151      # Show next steps
152      click.echo("\n" + "=" * 30)
153      click.echo("🚀 Next Steps:")
154      click.echo("=" * 30)
155  
156      # Only show cd if it's a different directory
157      if target_dir != current_dir:
158          click.echo(f"cd {target_dir}")
159  
160      click.echo("claude -p 'your prompt here'")
161  
162      if tracking_uri and tracking_uri.startswith("file://"):
163          click.echo("\n💡 View your traces:")
164          click.echo(f"   mlflow server --backend-store-uri {tracking_uri}")
165      elif not tracking_uri:
166          click.echo("\n💡 View your traces:")
167          click.echo("   mlflow server")
168      elif tracking_uri == "databricks":
169          click.echo("\n💡 View your traces in your Databricks workspace")
170  
171      click.echo("\n🔧 To disable tracing later:")
172      click.echo("   mlflow autolog claude --disable")
173  
174  
175  @claude.command("stop-hook", hidden=True)
176  def stop_hook() -> None:
177      """Hook handler invoked when a Claude Code conversation ends."""
178      stop_hook_handler()