|
| 1 | +""" |
| 2 | +Telemetry module for the promptfoo Python wrapper. |
| 3 | +
|
| 4 | +Sends anonymous usage analytics to PostHog to help improve promptfoo. |
| 5 | +Telemetry can be disabled by setting PROMPTFOO_DISABLE_TELEMETRY=1. |
| 6 | +""" |
| 7 | + |
| 8 | +import atexit |
| 9 | +import os |
| 10 | +import platform |
| 11 | +import sys |
| 12 | +import uuid |
| 13 | +from pathlib import Path |
| 14 | +from typing import Any, Optional |
| 15 | + |
| 16 | +import yaml |
| 17 | +from posthog import Posthog |
| 18 | + |
| 19 | +from . import __version__ |
| 20 | + |
| 21 | +# PostHog configuration - same as the main promptfoo TypeScript project. |
| 22 | +# NOTE: This is an intentionally public PostHog project API key: |
| 23 | +# - Safe to commit to source control (client-side telemetry key) |
| 24 | +# - Only allows sending anonymous usage events to the promptfoo PostHog project |
| 25 | +# - Does not grant administrative access to the PostHog account |
| 26 | +# - Abuse is mitigated by PostHog's built-in rate limiting |
| 27 | +# - Telemetry can be disabled via PROMPTFOO_DISABLE_TELEMETRY=1 |
| 28 | +_POSTHOG_HOST = "https://a.promptfoo.app" |
| 29 | +_POSTHOG_KEY = "phc_E5n5uHnDo2eREJL1uqX1cIlbkoRby4yFWt3V94HqRRg" |
| 30 | + |
| 31 | + |
| 32 | +def _get_env_bool(name: str) -> bool: |
| 33 | + """Check if an environment variable is set to a truthy value.""" |
| 34 | + value = os.environ.get(name, "").lower() |
| 35 | + return value in ("1", "true", "yes", "on") |
| 36 | + |
| 37 | + |
| 38 | +def _is_ci() -> bool: |
| 39 | + """Detect if running in a CI environment.""" |
| 40 | + ci_env_vars = [ |
| 41 | + "CI", |
| 42 | + "CONTINUOUS_INTEGRATION", |
| 43 | + "GITHUB_ACTIONS", |
| 44 | + "GITLAB_CI", |
| 45 | + "CIRCLECI", |
| 46 | + "TRAVIS", |
| 47 | + "JENKINS_URL", |
| 48 | + "BUILDKITE", |
| 49 | + "TEAMCITY_VERSION", |
| 50 | + "TF_BUILD", # Azure Pipelines |
| 51 | + ] |
| 52 | + return any(os.environ.get(var) for var in ci_env_vars) |
| 53 | + |
| 54 | + |
| 55 | +def _get_config_dir() -> Path: |
| 56 | + """Get the promptfoo config directory path.""" |
| 57 | + return Path.home() / ".promptfoo" |
| 58 | + |
| 59 | + |
| 60 | +def _read_global_config() -> dict[str, Any]: |
| 61 | + """Read the global promptfoo config from ~/.promptfoo/promptfoo.yaml.""" |
| 62 | + config_file = _get_config_dir() / "promptfoo.yaml" |
| 63 | + if config_file.exists(): |
| 64 | + try: |
| 65 | + with open(config_file) as f: |
| 66 | + config = yaml.safe_load(f) |
| 67 | + return config if isinstance(config, dict) else {} |
| 68 | + except Exception: |
| 69 | + return {} |
| 70 | + return {} |
| 71 | + |
| 72 | + |
| 73 | +def _write_global_config(config: dict[str, Any]) -> None: |
| 74 | + """Write the global promptfoo config to ~/.promptfoo/promptfoo.yaml.""" |
| 75 | + config_dir = _get_config_dir() |
| 76 | + config_dir.mkdir(parents=True, exist_ok=True) |
| 77 | + config_file = config_dir / "promptfoo.yaml" |
| 78 | + try: |
| 79 | + with open(config_file, "w") as f: |
| 80 | + yaml.dump(config, f, default_flow_style=False) |
| 81 | + except Exception: |
| 82 | + pass # Silently fail - telemetry should never break the CLI |
| 83 | + |
| 84 | + |
| 85 | +def _get_user_id() -> str: |
| 86 | + """Get or create a unique user ID stored in the global config.""" |
| 87 | + config = _read_global_config() |
| 88 | + user_id = config.get("id") |
| 89 | + |
| 90 | + if not user_id: |
| 91 | + user_id = str(uuid.uuid4()) |
| 92 | + config["id"] = user_id |
| 93 | + _write_global_config(config) |
| 94 | + |
| 95 | + return user_id |
| 96 | + |
| 97 | + |
| 98 | +def _get_user_email() -> Optional[str]: |
| 99 | + """Get the user email from the global config if set.""" |
| 100 | + config = _read_global_config() |
| 101 | + account = config.get("account", {}) |
| 102 | + return account.get("email") if isinstance(account, dict) else None |
| 103 | + |
| 104 | + |
| 105 | +class _Telemetry: |
| 106 | + """Internal telemetry client for the promptfoo Python wrapper.""" |
| 107 | + |
| 108 | + def __init__(self) -> None: |
| 109 | + self._client: Optional[Posthog] = None |
| 110 | + self._user_id: Optional[str] = None |
| 111 | + self._email: Optional[str] = None |
| 112 | + self._initialized = False |
| 113 | + |
| 114 | + @property |
| 115 | + def _disabled(self) -> bool: |
| 116 | + """Check if telemetry is disabled.""" |
| 117 | + return _get_env_bool("PROMPTFOO_DISABLE_TELEMETRY") or _get_env_bool("IS_TESTING") |
| 118 | + |
| 119 | + def _ensure_initialized(self) -> None: |
| 120 | + """Lazily initialize the telemetry client.""" |
| 121 | + if self._initialized: |
| 122 | + return |
| 123 | + |
| 124 | + self._initialized = True |
| 125 | + |
| 126 | + if self._disabled: |
| 127 | + return |
| 128 | + |
| 129 | + try: |
| 130 | + self._user_id = _get_user_id() |
| 131 | + self._email = _get_user_email() |
| 132 | + self._client = Posthog( |
| 133 | + project_api_key=_POSTHOG_KEY, |
| 134 | + host=_POSTHOG_HOST, |
| 135 | + ) |
| 136 | + except Exception: |
| 137 | + self._client = None # Silently fail |
| 138 | + |
| 139 | + def record(self, event_name: str, properties: Optional[dict[str, Any]] = None) -> None: |
| 140 | + """Record a telemetry event.""" |
| 141 | + if self._disabled: |
| 142 | + return |
| 143 | + |
| 144 | + self._ensure_initialized() |
| 145 | + |
| 146 | + if not self._client or not self._user_id: |
| 147 | + return |
| 148 | + |
| 149 | + try: |
| 150 | + enriched_properties: dict[str, Any] = { |
| 151 | + **(properties or {}), |
| 152 | + "packageVersion": __version__, |
| 153 | + "pythonVersion": platform.python_version(), |
| 154 | + "platform": sys.platform, |
| 155 | + "isRunningInCi": _is_ci(), |
| 156 | + "source": "python-wrapper", |
| 157 | + } |
| 158 | + |
| 159 | + # Only set email if present |
| 160 | + if self._email: |
| 161 | + enriched_properties["$set"] = {"email": self._email} |
| 162 | + |
| 163 | + self._client.capture( |
| 164 | + event=event_name, |
| 165 | + distinct_id=self._user_id, |
| 166 | + properties=enriched_properties, |
| 167 | + ) |
| 168 | + except Exception: |
| 169 | + pass # Silently fail - telemetry should never break the CLI |
| 170 | + |
| 171 | + def shutdown(self) -> None: |
| 172 | + """Shutdown the telemetry client and flush any pending events.""" |
| 173 | + if self._client: |
| 174 | + try: |
| 175 | + self._client.flush() |
| 176 | + self._client.shutdown() |
| 177 | + except Exception: |
| 178 | + pass # Silently fail |
| 179 | + finally: |
| 180 | + self._client = None |
| 181 | + |
| 182 | + |
| 183 | +# Global singleton instance |
| 184 | +_telemetry: Optional[_Telemetry] = None |
| 185 | + |
| 186 | + |
| 187 | +def _get_telemetry() -> _Telemetry: |
| 188 | + """Get the global telemetry instance.""" |
| 189 | + global _telemetry |
| 190 | + if _telemetry is None: |
| 191 | + _telemetry = _Telemetry() |
| 192 | + atexit.register(_telemetry.shutdown) |
| 193 | + return _telemetry |
| 194 | + |
| 195 | + |
| 196 | +def record_wrapper_used(method: str) -> None: |
| 197 | + """ |
| 198 | + Record that the Python wrapper was used. |
| 199 | +
|
| 200 | + Args: |
| 201 | + method: The execution method used - "global" for global promptfoo install, |
| 202 | + "npx" for npx fallback, or "error" if execution failed. |
| 203 | + """ |
| 204 | + _get_telemetry().record("wrapper_used", {"method": method, "wrapperType": "python"}) |
0 commit comments