1  """Configuration management for Claude Code integration with MLflow."""
  2  
  3  import json
  4  import os
  5  from dataclasses import dataclass
  6  from pathlib import Path
  7  from typing import Any
  8  
  9  from mlflow.environment_variables import (
 10      MLFLOW_EXPERIMENT_ID,
 11      MLFLOW_EXPERIMENT_NAME,
 12      MLFLOW_TRACKING_URI,
 13  )
 14  
 15  # Configuration field constants
 16  HOOK_FIELD_HOOKS = "hooks"
 17  HOOK_FIELD_COMMAND = "command"
 18  ENVIRONMENT_FIELD = "env"
 19  
 20  # MLflow environment variable constants
 21  MLFLOW_HOOK_IDENTIFIER = "mlflow autolog claude"
 22  # Legacy identifier used in older versions (inline python -c commands)
 23  MLFLOW_LEGACY_HOOK_IDENTIFIER = "mlflow.claude_code.hooks"
 24  MLFLOW_TRACING_ENABLED = "MLFLOW_CLAUDE_TRACING_ENABLED"
 25  
 26  
 27  @dataclass
 28  class TracingStatus:
 29      """Dataclass for tracing status information."""
 30  
 31      enabled: bool
 32      tracking_uri: str | None = None
 33      experiment_id: str | None = None
 34      experiment_name: str | None = None
 35      reason: str | None = None
 36  
 37  
 38  def load_claude_config(settings_path: Path) -> dict[str, Any]:
 39      """Load existing Claude configuration from settings file.
 40  
 41      Args:
 42          settings_path: Path to Claude settings.json file
 43  
 44      Returns:
 45          Configuration dictionary, empty dict if file doesn't exist or is invalid
 46      """
 47      if settings_path.exists():
 48          try:
 49              with open(settings_path, encoding="utf-8") as f:
 50                  return json.load(f)
 51          except (json.JSONDecodeError, IOError):
 52              return {}
 53      return {}
 54  
 55  
 56  def save_claude_config(settings_path: Path, config: dict[str, Any]) -> None:
 57      """Save Claude configuration to settings file.
 58  
 59      Args:
 60          settings_path: Path to Claude settings.json file
 61          config: Configuration dictionary to save
 62      """
 63      settings_path.parent.mkdir(parents=True, exist_ok=True)
 64      with open(settings_path, "w", encoding="utf-8") as f:
 65          json.dump(config, f, indent=2)
 66  
 67  
 68  def get_tracing_status(settings_path: Path) -> TracingStatus:
 69      """Get current tracing status from Claude settings.
 70  
 71      Args:
 72          settings_path: Path to Claude settings file
 73  
 74      Returns:
 75          TracingStatus with tracing status information
 76      """
 77      if not settings_path.exists():
 78          return TracingStatus(enabled=False, reason="No configuration found")
 79  
 80      config = load_claude_config(settings_path)
 81      env_vars = config.get(ENVIRONMENT_FIELD, {})
 82      enabled = env_vars.get(MLFLOW_TRACING_ENABLED) == "true"
 83  
 84      return TracingStatus(
 85          enabled=enabled,
 86          tracking_uri=env_vars.get(MLFLOW_TRACKING_URI.name),
 87          experiment_id=env_vars.get(MLFLOW_EXPERIMENT_ID.name),
 88          experiment_name=env_vars.get(MLFLOW_EXPERIMENT_NAME.name),
 89      )
 90  
 91  
 92  def get_env_var(var_name: str, default: str = "") -> str:
 93      """Get environment variable from Claude settings or OS environment as fallback.
 94  
 95      Project-specific configuration in settings.json takes precedence over
 96      global OS environment variables.
 97  
 98      Args:
 99          var_name: Environment variable name
100          default: Default value if not found anywhere
101  
102      Returns:
103          Environment variable value
104      """
105      # First check Claude settings (project-specific configuration takes priority)
106      try:
107          settings_path = Path(".claude/settings.json")
108          if settings_path.exists():
109              config = load_claude_config(settings_path)
110              env_vars = config.get(ENVIRONMENT_FIELD, {})
111              value = env_vars.get(var_name)
112              if value is not None:
113                  return value
114      except Exception:
115          pass
116  
117      # Fallback to OS environment
118      value = os.environ.get(var_name)
119      if value is not None:
120          return value
121  
122      return default
123  
124  
125  def setup_environment_config(
126      settings_path: Path,
127      tracking_uri: str | None = None,
128      experiment_id: str | None = None,
129      experiment_name: str | None = None,
130  ) -> None:
131      """Set up MLflow environment variables in Claude settings.
132  
133      Args:
134          settings_path: Path to Claude settings file
135          tracking_uri: MLflow tracking URI, defaults to local file storage
136          experiment_id: MLflow experiment ID (takes precedence over name)
137          experiment_name: MLflow experiment name
138      """
139      config = load_claude_config(settings_path)
140  
141      if ENVIRONMENT_FIELD not in config:
142          config[ENVIRONMENT_FIELD] = {}
143  
144      # Always enable tracing
145      config[ENVIRONMENT_FIELD][MLFLOW_TRACING_ENABLED] = "true"
146  
147      # Set tracking URI
148      if tracking_uri:
149          config[ENVIRONMENT_FIELD][MLFLOW_TRACKING_URI.name] = tracking_uri
150  
151      # Set experiment configuration (ID takes precedence over name)
152      if experiment_id:
153          config[ENVIRONMENT_FIELD][MLFLOW_EXPERIMENT_ID.name] = experiment_id
154          config[ENVIRONMENT_FIELD].pop(MLFLOW_EXPERIMENT_NAME.name, None)
155      elif experiment_name:
156          config[ENVIRONMENT_FIELD][MLFLOW_EXPERIMENT_NAME.name] = experiment_name
157          config[ENVIRONMENT_FIELD].pop(MLFLOW_EXPERIMENT_ID.name, None)
158  
159      save_claude_config(settings_path, config)