done

Author: shreyaspimpalgaonkar

poetry.lock

@@ -37,6 +37,7 @@ vertexai= "1.71.1"
 vllm = { version = "^0.6.3", optional = true }
 aiodocker = { version = "^0.24.0", optional = true }
 ray = { version = "^2.41.0", optional = true }
+posthog = "^3.11.0"
 
 [tool.poetry.extras]
 vllm = ["vllm"]

pyproject.toml

@@ -1,3 +1,4 @@
 """Defaults used across curator."""
 
 BATCH_REQUEST_ID_TAG = "custom_id"
+_CURATOR_DEFAULT_CACHE_DIR = "~/.cache/curator"

src/bespokelabs/curator/constants.py

@@ -11,6 +11,7 @@
 from pydantic import BaseModel
 from xxhash import xxh64
 
+from bespokelabs.curator.constants import _CURATOR_DEFAULT_CACHE_DIR
 from bespokelabs.curator.db import MetadataDB
 from bespokelabs.curator.llm.prompt_formatter import PromptFormatter
 from bespokelabs.curator.request_processor._factory import _RequestProcessorFactory
@@ -19,7 +20,6 @@
 if TYPE_CHECKING:
     from datasets import Dataset
 
-_CURATOR_DEFAULT_CACHE_DIR = "~/.cache/curator"
 T = TypeVar("T")
 _DictOrBaseModel = Dict[str, Any] | BaseModel
 logger = logging.getLogger(__name__)

src/bespokelabs/curator/llm/llm.py

@@ -10,6 +10,7 @@
 from rich.progress import BarColumn, Progress, TextColumn, TimeElapsedColumn
 from rich.table import Table
 
+from bespokelabs.curator.telemetry.client import TelemetryEvent, telemetry_client
 from bespokelabs.curator.types.generic_batch import GenericBatch, GenericBatchStatus
 from bespokelabs.curator.types.generic_response import TokenUsage
 
@@ -92,6 +93,14 @@ def stop_tracker(self):
             self._progress.stop()
             self.display_final_stats()
 
+        # update anonymized telemetry
+        telemetry_client.capture(
+            TelemetryEvent(
+                event_type="BatchRequest",
+                metadata=self.model_dump(),
+            )
+        )
+
     def update_display(self):
         """Update statistics with token usage and cost information."""
         # Format display texts

src/bespokelabs/curator/status_tracker/batch_status_tracker.py

@@ -1,7 +1,7 @@
 import logging
 import time
 import typing as t
-from dataclasses import dataclass, field
+from dataclasses import asdict, dataclass, field
 from enum import Enum
 from typing import Optional
 
@@ -13,18 +13,23 @@
 from rich.progress import BarColumn, Progress, TextColumn, TimeElapsedColumn, TimeRemainingColumn
 from rich.table import Table
 
+from bespokelabs.curator.telemetry.client import TelemetryEvent, telemetry_client
 from bespokelabs.curator.types.generic_response import TokenUsage
 
 logger = logging.getLogger(__name__)
 
 
-class TokenLimitStrategy(Enum):
+class TokenLimitStrategy(str, Enum):
     """Token limit Strategy enum."""
 
     combined = "combined"
     seperate = "seperate"
     default = "combined"
 
+    def __str__(self):
+        """String representation of the token limit strategy."""
+        return self.value
+
 
 class _TokenCount(BaseModel):
     input: int | float | None = 0
@@ -279,6 +284,13 @@ def stop_tracker(self):
 
         self._console.print(table)
 
+        telemetry_client.capture(
+            TelemetryEvent(
+                event_type="OnlineRequest",
+                metadata=asdict(self),
+            )
+        )
+
     def __str__(self):
         """String representation of the status tracker."""
         return (

src/bespokelabs/curator/status_tracker/online_status_tracker.py

@@ -1,35 +1,58 @@
 import logging
 import os
+import uuid
 from dataclasses import dataclass
 from typing import Optional
 
 import posthog
-from bespokelabs.curator.telemetry.events import TelemetryEvent
+
+from bespokelabs.curator.constants import _CURATOR_DEFAULT_CACHE_DIR
 
 logger = logging.getLogger(__name__)
 
+
+def _random_distinct_id():
+    # check if cache
+    curator_cache_dir = os.environ.get(
+        "CURATOR_CACHE_DIR",
+        os.path.expanduser(_CURATOR_DEFAULT_CACHE_DIR),
+    )
+    os.makedirs(curator_cache_dir, exist_ok=True)
+    distinct_id_file = os.path.join(curator_cache_dir, ".distinct_id.txt")
+
+    if os.path.exists(distinct_id_file):
+        with open(distinct_id_file) as f:
+            random_distinct_id = uuid.UUID(f.read().strip())
+    else:
+        random_distinct_id = uuid.uuid4()
+        with open(distinct_id_file, "w") as f:
+            f.write(str(random_distinct_id))
+
+    return random_distinct_id
+
+
 @dataclass
 class TelemetryEvent:
     """Base class for all telemetry events."""
+
     event_type: str
-    event_id: str
     metadata: dict
+    distinct_id: str = _random_distinct_id()
 
 
 @dataclass
 class PosthogConfig:
     """Configuration settings for PostHog client."""
+
     api_key: str
     enabled: bool = True
     debug: bool = False
-    disable_geoip: bool = False
     host: Optional[str] = None
 
 
 class PosthogClient:
-    """
-    Client for sending telemetry events to PostHog analytics.
-    
+    """Client for sending telemetry events to PostHog analytics.
+
     This uses a write-only project API key that can only create new events.
     It cannot read events or access other PostHog data, making it safe for public apps.
     """
@@ -40,42 +63,34 @@ def __init__(self, config: PosthogConfig):
         self._initialize_client()
 
     def _initialize_client(self) -> None:
-        """Set up the PostHog client with configuration settings"""
-        
+        """Set up the PostHog client with configuration settings."""
         posthog.project_api_key = self.config.api_key
         posthog.debug = self.config.debug
-        posthog.disable_geoip = self.config.disable_geoip
-        
+
         if self.config.host:
             posthog.host = self.config.host
 
     def capture(self, event: TelemetryEvent) -> None:
-        """
-        Capture and send a telemetry event to PostHog
-        
+        """Capture and send a telemetry event to PostHog.
+
         Args:
             event: The telemetry event to capture
         """
         if not self.config.enabled:
             return
-            
+
         try:
-            posthog.capture(
-                distinct_id=event.event_id,
-                event=event.event_type,
-                properties=event.properties
-            )
-        except Exception as e:
-            logger.error(f"Failed to capture telemetry event: {e}")
+            posthog.capture(distinct_id=event.distinct_id, event=event.event_type, properties=event.metadata)
+        except Exception:
+            pass
 
 
 # Initialize the telemetry client with environment-based configuration
 config = PosthogConfig(
-    api_key=os.getenv("POSTHOG_API_KEY"),
+    api_key="phc_HGGTf1LmtsUnBaVBufgIwRsAwdkvH3cSsDKgW5RnJz8",
     enabled=os.getenv("TELEMETRY_ENABLED", "true").lower() in ("true", "1", "t"),
     debug=os.getenv("DEBUG_MODE", "false").lower() in ("true", "1", "t"),
-    disable_geoip=os.getenv("TELEMETRY_ENABLED", "false").lower() in ("true", "1", "t"),
-    host=os.getenv("POSTHOG_HOST")
+    host=os.getenv("POSTHOG_HOST"),
 )
 
 telemetry_client = PosthogClient(config=config)