Response parser and split up class diagrams

sia/system_metrics.py

@@ -0,0 +1,122 @@
+import threading
+import time
+from typing import Optional
+import xml.etree.ElementTree as ET
+import psutil
+
+class SystemMetrics:
+    """
+    Tracks system metrics including CPU, GPU, memory and disk usage.
+    Maintains rolling averages for CPU and GPU between context generations.
+    
+    Attributes:
+        _stop_event: Threading event to signal monitor thread to stop
+        _monitor_thread: Background thread that collects usage data
+        _cpu_samples: List of CPU usage samples since last context
+        _gpu_samples: List of GPU usage samples since last context
+        _metrics_lock: Lock for thread-safe access to samples
+    """
+    
+    def __init__(self, sample_interval: float = 0.1):
+        """
+        Initialize system metrics monitoring.
+        
+        Args:
+            sample_interval: Seconds between usage samples
+        """
+        self._stop_event = threading.Event()
+        self._cpu_samples = []
+        self._gpu_samples = []
+        self._metrics_lock = threading.Lock()
+        self._sample_interval = sample_interval
+        
+        # Start monitoring thread
+        self._monitor_thread = threading.Thread(
+            target=self._monitor_loop,
+            daemon=True
+        )
+        self._monitor_thread.start()
+        
+    def _get_gpu_usage(self) -> Optional[float]:
+        """
+        Get current GPU usage percentage.
+        Returns None if GPU monitoring not available.
+        """
+        try:
+            import pynvml
+            pynvml.nvmlInit()
+            handle = pynvml.nvmlDeviceGetHandleByIndex(0)
+            info = pynvml.nvmlDeviceGetUtilizationRates(handle)
+            pynvml.nvmlShutdown()
+            return info.gpu
+        except:
+            return None
+
+    def _monitor_loop(self):
+        """
+        Background loop that collects system usage samples.
+        Runs until stop event is set.
+        """
+        while not self._stop_event.is_set():
+            with self._metrics_lock:
+                # Get CPU usage across all cores
+                self._cpu_samples.append(psutil.cpu_percent())
+                
+                # Get GPU usage if available
+                gpu_usage = self._get_gpu_usage()
+                if gpu_usage is not None:
+                    self._gpu_samples.append(gpu_usage)
+                    
+            time.sleep(self._sample_interval)
+
+    def generate_context(self, context_usage: float) -> ET.Element:
+        """
+        Generate XML context element with current system metrics.
+        Clears usage samples after generating averages.
+        
+        Args:
+            context_usage: Current context size usage (0-1)
+            
+        Returns:
+            ET.Element: Context element with system metrics
+        """
+        # Create context element
+        context = ET.Element("context")
+        
+        # Add timestamp
+        context.set("time", time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()))
+        
+        # Add usage metrics
+        with self._metrics_lock:
+            # CPU average (or 0 if no samples)
+            cpu_avg = round(sum(self._cpu_samples) / len(self._cpu_samples)) if self._cpu_samples else 0
+            context.set("cpu", str(cpu_avg))
+            self._cpu_samples.clear()
+            
+            # GPU average (or 0 if no samples)
+            gpu_avg = round(sum(self._gpu_samples) / len(self._gpu_samples)) if self._gpu_samples else 0
+            context.set("gpu", str(gpu_avg))
+            self._gpu_samples.clear()
+        
+        # Memory usage in bytes
+        memory = psutil.virtual_memory()
+        context.set("memory_used", str(memory.used))
+        context.set("memory_total", str(memory.total))
+        
+        # Disk usage in bytes (root partition)
+        disk = psutil.disk_usage('/')
+        context.set("disk_used", str(disk.used))
+        context.set("disk_total", str(disk.total))
+        
+        # Context usage (0-100)
+        context.set("context", str(round(context_usage * 100)))
+        
+        # Standard input buffer size 
+        context.set("stdin", "0")  # Set by agent
+        
+        return context
+        
+    def stop(self):
+        """Stop the monitoring thread and clean up."""
+        self._stop_event.set()
+        self._monitor_thread.join()