Merge pull request #336 from jamesmurdza/fix/refactor-python-telemetry

Refactor Python telemetry library

libs/python/agent/agent/__init__.py

@@ -32,9 +32,6 @@ try:
         record_event,
     )
 
-    # Import set_dimension from our own telemetry module
-    from .telemetry import set_dimension
-
     # Check if telemetry is enabled
     if is_telemetry_enabled():
         logger.info("Telemetry is enabled")
@@ -49,9 +46,6 @@ try:
             },
         )
 
-        # Set the package version as a dimension
-        set_dimension("agent_version", __version__)
-
         # Flush events to ensure they're sent
         flush()
     else:

libs/python/agent/agent/callbacks/telemetry.py

@@ -7,13 +7,18 @@ import uuid
 from typing import List, Dict, Any, Optional, Union
 
 from .base import AsyncCallbackHandler
-from ..telemetry import (
+from core.telemetry import (
     record_event,
     is_telemetry_enabled,
-    set_dimension,
-    SYSTEM_INFO,
 )
 
+import platform
+
+SYSTEM_INFO = {
+    "os": platform.system().lower(),
+    "os_version": platform.release(),
+    "python_version": platform.python_version(),
+}
 
 class TelemetryCallback(AsyncCallbackHandler):
     """
@@ -65,11 +70,6 @@ class TelemetryCallback(AsyncCallbackHandler):
             **SYSTEM_INFO
         }
         
-        # Set session-level dimensions
-        set_dimension("session_id", self.session_id)
-        set_dimension("agent_type", agent_info["agent_type"])
-        set_dimension("model", agent_info["model"])
-        
         record_event("agent_session_start", agent_info)
     
     async def on_run_start(self, kwargs: Dict[str, Any], old_items: List[Dict[str, Any]]) -> None:
@@ -98,7 +98,6 @@ class TelemetryCallback(AsyncCallbackHandler):
             if trajectory:
                 run_data["uploaded_trajectory"] = trajectory
         
-        set_dimension("run_id", self.run_id)
         record_event("agent_run_start", run_data)
     
     async def on_run_end(self, kwargs: Dict[str, Any], old_items: List[Dict[str, Any]], new_items: List[Dict[str, Any]]) -> None:

libs/python/agent/agent/telemetry.py

@@ -1,142 +0,0 @@
-"""Agent telemetry for tracking anonymous usage and feature usage."""
-
-import logging
-import os
-import platform
-import sys
-from typing import Dict, Any, Callable
-
-# Import the core telemetry module
-TELEMETRY_AVAILABLE = False
-
-
-# Local fallbacks in case core telemetry isn't available
-def _noop(*args: Any, **kwargs: Any) -> None:
-    """No-op function for when telemetry is not available."""
-    pass
-
-
-# Define default functions with unique names to avoid shadowing
-_default_record_event = _noop
-_default_increment_counter = _noop
-_default_set_dimension = _noop
-_default_get_telemetry_client = lambda: None
-_default_flush = _noop
-_default_is_telemetry_enabled = lambda: False
-_default_is_telemetry_globally_disabled = lambda: True
-
-# Set the actual functions to the defaults initially
-record_event = _default_record_event
-increment_counter = _default_increment_counter
-set_dimension = _default_set_dimension
-get_telemetry_client = _default_get_telemetry_client
-flush = _default_flush
-is_telemetry_enabled = _default_is_telemetry_enabled
-is_telemetry_globally_disabled = _default_is_telemetry_globally_disabled
-
-logger = logging.getLogger("agent.telemetry")
-
-try:
-    # Import from core telemetry
-    from core.telemetry import (
-        record_event as core_record_event,
-        increment as core_increment,
-        get_telemetry_client as core_get_telemetry_client,
-        flush as core_flush,
-        is_telemetry_enabled as core_is_telemetry_enabled,
-        is_telemetry_globally_disabled as core_is_telemetry_globally_disabled,
-    )
-
-    # Override the default functions with actual implementations
-    record_event = core_record_event
-    get_telemetry_client = core_get_telemetry_client
-    flush = core_flush
-    is_telemetry_enabled = core_is_telemetry_enabled
-    is_telemetry_globally_disabled = core_is_telemetry_globally_disabled
-
-    def increment_counter(counter_name: str, value: int = 1) -> None:
-        """Wrapper for increment to maintain backward compatibility."""
-        if is_telemetry_enabled():
-            core_increment(counter_name, value)
-
-    def set_dimension(name: str, value: Any) -> None:
-        """Set a dimension that will be attached to all events."""
-        logger.debug(f"Setting dimension {name}={value}")
-
-    TELEMETRY_AVAILABLE = True
-    logger.info("Successfully imported telemetry")
-except ImportError as e:
-    logger.warning(f"Could not import telemetry: {e}")
-    logger.debug("Telemetry not available, using no-op functions")
-
-# Get system info once to use in telemetry
-SYSTEM_INFO = {
-    "os": platform.system().lower(),
-    "os_version": platform.release(),
-    "python_version": platform.python_version(),
-}
-
-
-def enable_telemetry() -> bool:
-    """Enable telemetry if available.
-
-    Returns:
-        bool: True if telemetry was successfully enabled, False otherwise
-    """
-    global TELEMETRY_AVAILABLE, record_event, increment_counter, get_telemetry_client, flush, is_telemetry_enabled, is_telemetry_globally_disabled
-
-    # Check if globally disabled using core function
-    if TELEMETRY_AVAILABLE and is_telemetry_globally_disabled():
-        logger.info("Telemetry is globally disabled via environment variable - cannot enable")
-        return False
-
-    # Already enabled
-    if TELEMETRY_AVAILABLE:
-        return True
-
-    # Try to import and enable
-    try:
-        from core.telemetry import (
-            record_event,
-            increment,
-            get_telemetry_client,
-            flush,
-            is_telemetry_globally_disabled,
-        )
-
-        # Check again after import
-        if is_telemetry_globally_disabled():
-            logger.info("Telemetry is globally disabled via environment variable - cannot enable")
-            return False
-
-        TELEMETRY_AVAILABLE = True
-        logger.info("Telemetry successfully enabled")
-        return True
-    except ImportError as e:
-        logger.warning(f"Could not enable telemetry: {e}")
-        return False
-
-
-def is_telemetry_enabled() -> bool:
-    """Check if telemetry is enabled.
-
-    Returns:
-        bool: True if telemetry is enabled, False otherwise
-    """
-    # Use the core function if available, otherwise use our local flag
-    if TELEMETRY_AVAILABLE:
-        from core.telemetry import is_telemetry_enabled as core_is_enabled
-
-        return core_is_enabled()
-    return False
-
-
-def record_agent_initialization() -> None:
-    """Record when an agent instance is initialized."""
-    if TELEMETRY_AVAILABLE and is_telemetry_enabled():
-        record_event("agent_initialized", SYSTEM_INFO)
-
-        # Set dimensions that will be attached to all events
-        set_dimension("os", SYSTEM_INFO["os"])
-        set_dimension("os_version", SYSTEM_INFO["os_version"])
-        set_dimension("python_version", SYSTEM_INFO["python_version"])