diff --git a/src/axolotl/telemetry/manager.py b/src/axolotl/telemetry/manager.py index 646688c73..06699ebff 100644 --- a/src/axolotl/telemetry/manager.py +++ b/src/axolotl/telemetry/manager.py @@ -165,10 +165,10 @@ class TelemetryManager: if not properties: return {} - # TODO: Keep this up to date with any config schema changes + # NOTE: Keep this up to date with any config schema changes path_indicators = {"path", "dir"} - def redact_value(value: str, key: Any = None) -> Any: + def redact_value(value: Any, key: str = "") -> Any: """Recursively sanitize values, redacting those with path-like keys""" # Special case: base_model should be redacted if org is not whitelisted if key == "base_model": @@ -181,11 +181,11 @@ class TelemetryManager: if any(indicator in key.lower() for indicator in path_indicators): return "[REDACTED]" - # Handle nested structures - if isinstance(value, dict): - return {k: redact_value(v, k) for k, v in value.items()} - if isinstance(value, list): - return [redact_value(item) for item in value] + # Handle nested structures + if isinstance(value, dict): + return {k: redact_value(v, k) for k, v in value.items()} + if isinstance(value, list): + return [redact_value(item) for item in value] return value diff --git a/tests/telemetry/test_manager.py b/tests/telemetry/test_manager.py index 0085f7b0f..4608b337c 100644 --- a/tests/telemetry/test_manager.py +++ b/tests/telemetry/test_manager.py @@ -196,8 +196,9 @@ def test_redacted_properties(manager): "path_to_model": "models/llama/7b", "message": "Training started", # Should not be redacted "metrics": {"loss": 0.5, "accuracy": 0.95}, # Should not be redacted + "base_model": "models/local_model", "nested": { - "model_path": "/models/local/weights.pt", + "model_path": "/models/my_model", "root_dir": "/home/user/projects", "stats": {"steps": 1000, "epochs": 3}, # Should not be redacted }, @@ -211,10 +212,11 @@ def test_redacted_properties(manager): # Get the sanitized properties that were sent sanitized = mock_capture.call_args[1]["properties"] - # Check that path-like keys were redacted + # Check that path-like and base_model keys were redacted assert sanitized["filepath"] == "[REDACTED]" assert sanitized["windows_path"] == "[REDACTED]" assert sanitized["path_to_model"] == "[REDACTED]" + assert sanitized["base_model"] == "[REDACTED]" # Check that non-path values were preserved assert sanitized["message"] == "Training started"