telemetry with posthog

Author: nithinjamespallicka

pyproject.toml

@@ -38,6 +38,7 @@ dependencies = [
     "fastmcp>=2.10.5",
     "mcp-use[search]>=1.3.7",
     "onnxruntime==1.19.2 ; sys_platform == 'darwin' and platform_machine == 'x86_64'",
+    "posthog>=6.5.0",
 ]
 
 [project.urls]

src/datu/app_config.py

@@ -17,6 +17,7 @@
 from datu.integrations.config import IntegrationConfigs
 from datu.mcp.config import MCPConfig
 from datu.services.config import SchemaRAGConfig
+from datu.telemetry.config import TelemetryConfig
 
 
 class Environment(Enum):
@@ -110,7 +111,8 @@ class DatuConfig(BaseSettings):
         description="Configuration settings for schema RAG (Retrieval-Augmented Generation).",
     )
     enable_anonymization: bool = False
-
+    enable_anonymized_telemetry: bool = True
+    telemetry: TelemetryConfig | None = Field(default_factory=TelemetryConfig)
     model_config = SettingsConfigDict(
         env_prefix="datu_",
         env_nested_delimiter="__",

src/datu/main.py

@@ -13,6 +13,8 @@
 from datu.app_config import get_logger, settings
 from datu.routers import chat, metadata, transformations
 from datu.schema_extractor.schema_cache import load_schema_cache
+from datu.telemetry.product.events import OpenAIEvents
+from datu.telemetry.product.posthog import get_posthog_client
 
 logger = get_logger(__name__)
 
@@ -56,6 +58,8 @@ def start_app() -> None:
     It also sets the logging level based on the configuration settings.
     """
     logger.info("Starting the FastAPI application...")
+    posthog_client = get_posthog_client()
+    posthog_client.capture(OpenAIEvents())
     uvicorn.run(app, host=settings.host, port=settings.port)
 
 

src/datu/telemetry/__init__.py

@@ -0,0 +1,19 @@
+"""Telemetry configuration settings."""
+
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+
+class TelemetryConfig(BaseSettings):
+    """Telemetry configuration settings."""
+
+    api_key: str = "phc_m74dfR9nLpm2nipvkL2swyFDtNuQNC9o2FL2CSbh6Je"
+    package_name: str = "datu-core"
+
+    model_config = SettingsConfigDict(
+        env_prefix="telemetry_",
+        env_nested_delimiter="__",
+    )
+
+
+def get_telemetry_settings() -> TelemetryConfig:
+    return TelemetryConfig()

src/datu/telemetry/config.py

@@ -0,0 +1,50 @@
+from typing import Any, ClassVar, Dict
+
+from datu.app_config import get_app_settings, get_logger
+
+app_settings = get_app_settings()
+
+
+class ProductTelemetryEvent:
+    """Base class for all telemetry events."""
+
+    max_batch_size: ClassVar[int] = 1
+
+    def __init__(self, **kwargs):
+        self._props = kwargs
+        self.batch_size = 1
+
+    @property
+    def name(self) -> str:
+        return self.__class__.__name__
+
+    @property
+    def properties(self) -> Dict[str, Any]:
+        return self._props
+
+    @property
+    def batch_key(self) -> str:
+        return self.name
+
+    def batch(self, other: "ProductTelemetryEvent") -> "ProductTelemetryEvent":
+        """Simple batch: append counts together."""
+        if self.name != other.name:
+            raise ValueError("Cannot batch different event types")
+        self.batch_size += other.batch_size
+        return self
+
+
+class MCPClientStartEvent(ProductTelemetryEvent):
+    """Event for when the MCP client starts."""
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self._props["mcp_client_enabled"] = app_settings.mcp_client_enabled
+
+
+class OpenAIEvents(ProductTelemetryEvent):
+    """Event for OpenAI-related telemetry."""
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self._props["openai_model"] = app_settings.openai_model

src/datu/telemetry/product/__init__.py

@@ -0,0 +1,109 @@
+import importlib
+import logging
+import platform
+import sys
+import uuid
+from functools import lru_cache
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+import posthog
+
+from datu.app_config import get_app_settings, get_logger
+from datu.telemetry.config import TelemetryConfig as TelemetrySettings
+from datu.telemetry.product.events import ProductTelemetryEvent
+
+app_settings = get_app_settings()
+logger = get_logger(__name__)
+
+POSTHOG_EVENT_SETTINGS = {"$process_person_profile": False}
+
+
+class PostHogClient:
+    """Telemetry client with basic batching + config via Pydantic."""
+
+    UNKNOWN_USER_ID = "UNKNOWN"
+    USER_ID_PATH = Path.home() / ".cache" / "datu-core" / "telemetry_user_id"
+
+    def __init__(self, settings: Optional[TelemetrySettings]) -> None:
+        self.settings = settings or TelemetrySettings()
+        self._batched_events: Dict[str, ProductTelemetryEvent] = {}
+        self._user_id: str = ""
+        self._user_id_path: Path = self.USER_ID_PATH
+
+        if not app_settings.enable_anonymized_telemetry or "pytest" in sys.modules:
+            posthog.disabled = True
+        else:
+            logger.info("Enabled anonymized telemetry. See https://docs.datu.fi for more information.")
+        posthog.api_key = self.settings.api_key
+        posthog_logger = logging.getLogger("posthog")
+        posthog_logger.disabled = True
+
+    @property
+    def user_id(self) -> str:
+        if self._user_id:
+            return self._user_id
+
+        try:
+            if not self._user_id_path.exists():
+                self._user_id_path.parent.mkdir(parents=True, exist_ok=True)
+                new_id = str(uuid.uuid4())
+                self._user_id_path.write_text(new_id)
+                self._user_id = new_id
+            else:
+                self._user_id = self._user_id_path.read_text().strip()
+        except Exception:
+            self._user_id = self.UNKNOWN_USER_ID
+
+        return self._user_id
+
+    def _base_context(self) -> Dict[str, Any]:
+        try:
+            pkg_version = importlib.metadata.version(self.settings.package_name)
+        except importlib.metadata.PackageNotFoundError:
+            pkg_version = "unknown"
+
+        return {
+            "python_version": sys.version.split()[0],
+            "os": platform.system(),
+            "os_version": platform.release(),
+            "package_version": pkg_version,
+        }
+
+    def capture(self, event: ProductTelemetryEvent) -> None:
+        """Capture an event (with simple batching)."""
+        if not app_settings.enable_anonymized_telemetry or not self.settings.api_key:
+            return
+
+        if event.max_batch_size == 1:
+            self._send(event)
+            return
+
+        batch_key = event.batch_key
+        if batch_key not in self._batched_events:
+            self._batched_events[batch_key] = event
+            return
+
+        batched = self._batched_events[batch_key].batch(event)
+        self._batched_events[batch_key] = batched
+
+        if batched.batch_size >= batched.max_batch_size:
+            self._send(batched)
+            del self._batched_events[batch_key]
+
+    def _send(self, event: ProductTelemetryEvent) -> None:
+        try:
+            posthog.capture(
+                distinct_id=self.user_id,
+                event=event.name,
+                properties={**self._base_context(), **POSTHOG_EVENT_SETTINGS, **event.properties},
+            )
+        except Exception:
+            # silently fail
+            pass
+
+
+@lru_cache(maxsize=1)
+def get_posthog_client() -> PostHogClient:
+    """Get the PostHog telemetry client."""
+    return PostHogClient(app_settings.telemetry)