Coverage for src/chuck_data/metrics_collector.py: 0%

51 statements  

« prev     ^ index     » next       coverage.py v7.8.0, created at 2025-06-05 22:56 -0700

1""" 

2Metrics collection service for tracking usage events. 

3 

4This module provides functionality to collect and send metrics about usage 

5of the application to help improve its features and performance. 

6""" 

7 

8import json 

9import logging 

10from typing import Any, Dict, List, Optional, Union 

11from .clients.amperity import AmperityAPIClient 

12 

13from .config import get_config_manager, get_amperity_token 

14 

15 

16class MetricsCollector: 

17 """Collects and sends usage metrics to the Amperity API.""" 

18 

19 def __init__(self): 

20 """Initialize the metrics collector.""" 

21 self.config_manager = get_config_manager() 

22 self._client = AmperityAPIClient() 

23 

24 def _should_track(self) -> bool: 

25 """ 

26 Determine if metrics should be tracked based on user consent. 

27 

28 Returns: 

29 bool: True if user has provided consent, False otherwise. 

30 """ 

31 return self.config_manager.get_config().usage_tracking_consent 

32 

33 def _get_chuck_configuration_for_metric(self) -> Dict[str, Any]: 

34 """ 

35 Get the configuration settings relevant for metrics. 

36 

37 Returns: 

38 Dict[str, Any]: Dictionary of configuration values. 

39 """ 

40 config = self.config_manager.get_config() 

41 return { 

42 "workspace_url": config.workspace_url, 

43 "active_catalog": config.active_catalog, 

44 "active_schema": config.active_schema, 

45 "active_model": config.active_model, 

46 } 

47 

48 def track_event( 

49 self, 

50 prompt: Optional[str] = None, 

51 tools: Optional[Union[List[Dict[str, Any]], Dict[str, Any]]] = None, 

52 conversation_history: Optional[List[Dict[str, Any]]] = None, 

53 error: Optional[str] = None, 

54 additional_data: Optional[Dict[str, Any]] = None, 

55 ) -> bool: 

56 """ 

57 Track a usage event with provided data. 

58 

59 Args: 

60 prompt: The user prompt or query that triggered this event 

61 tools: Tool usage information for this event 

62 conversation_history: Previous conversation messages 

63 error: Error information if this is an error event 

64 additional_data: Any additional context-specific data 

65 

66 Returns: 

67 bool: True if the metrics were sent successfully, False otherwise. 

68 """ 

69 if not self._should_track(): 

70 logging.debug("Metrics tracking skipped - user has not provided consent") 

71 return False 

72 

73 try: 

74 # Convert tools to list if it's a dict 

75 if tools and isinstance(tools, dict): 

76 tools = [tools] 

77 

78 # Prepare the payload 

79 payload = { 

80 "event": "USAGE", # All events are USAGE events 

81 "configuration": self._get_chuck_configuration_for_metric(), 

82 } 

83 

84 # Add optional fields if provided 

85 if prompt: 

86 payload["prompt"] = prompt 

87 if tools: 

88 payload["tools"] = tools 

89 if conversation_history: 

90 payload["conversation_history"] = conversation_history 

91 if error: 

92 payload["error"] = error 

93 if additional_data: 

94 payload["additional_data"] = additional_data 

95 

96 # Send the metric 

97 return self.send_metric(payload) 

98 except Exception as e: 

99 logging.debug(f"Error tracking metrics: {e}", exc_info=True) 

100 return False 

101 

102 def send_metric(self, payload: Dict[str, Any]) -> bool: 

103 """ 

104 Send the collected metric to the Amperity API. 

105 

106 Args: 

107 payload: The data payload to send 

108 

109 Returns: 

110 bool: True if sent successfully, False otherwise. 

111 """ 

112 try: 

113 

114 token = get_amperity_token() 

115 if not token: 

116 logging.debug("Cannot send metrics - no authentication token available") 

117 return False 

118 

119 # Convert the payload to a JSON string for logging 

120 payload_str = json.dumps(payload) 

121 logging.debug(f"Sending metric: {payload_str[:100]}...") 

122 

123 return self._client.submit_metrics(payload, token) 

124 except Exception as e: 

125 logging.debug(f"Error sending metrics: {e}", exc_info=True) 

126 return False 

127 

128 

129# Global metrics collector instance 

130_metrics_collector = MetricsCollector() 

131 

132 

133def get_metrics_collector() -> MetricsCollector: 

134 """ 

135 Get the global metrics collector instance. 

136 

137 Returns: 

138 MetricsCollector: The global metrics collector instance. 

139 """ 

140 return _metrics_collector