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
« prev ^ index » next coverage.py v7.8.0, created at 2025-06-05 22:56 -0700
1"""
2Metrics collection service for tracking usage events.
4This module provides functionality to collect and send metrics about usage
5of the application to help improve its features and performance.
6"""
8import json
9import logging
10from typing import Any, Dict, List, Optional, Union
11from .clients.amperity import AmperityAPIClient
13from .config import get_config_manager, get_amperity_token
16class MetricsCollector:
17 """Collects and sends usage metrics to the Amperity API."""
19 def __init__(self):
20 """Initialize the metrics collector."""
21 self.config_manager = get_config_manager()
22 self._client = AmperityAPIClient()
24 def _should_track(self) -> bool:
25 """
26 Determine if metrics should be tracked based on user consent.
28 Returns:
29 bool: True if user has provided consent, False otherwise.
30 """
31 return self.config_manager.get_config().usage_tracking_consent
33 def _get_chuck_configuration_for_metric(self) -> Dict[str, Any]:
34 """
35 Get the configuration settings relevant for metrics.
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 }
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.
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
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
73 try:
74 # Convert tools to list if it's a dict
75 if tools and isinstance(tools, dict):
76 tools = [tools]
78 # Prepare the payload
79 payload = {
80 "event": "USAGE", # All events are USAGE events
81 "configuration": self._get_chuck_configuration_for_metric(),
82 }
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
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
102 def send_metric(self, payload: Dict[str, Any]) -> bool:
103 """
104 Send the collected metric to the Amperity API.
106 Args:
107 payload: The data payload to send
109 Returns:
110 bool: True if sent successfully, False otherwise.
111 """
112 try:
114 token = get_amperity_token()
115 if not token:
116 logging.debug("Cannot send metrics - no authentication token available")
117 return False
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]}...")
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
129# Global metrics collector instance
130_metrics_collector = MetricsCollector()
133def get_metrics_collector() -> MetricsCollector:
134 """
135 Get the global metrics collector instance.
137 Returns:
138 MetricsCollector: The global metrics collector instance.
139 """
140 return _metrics_collector