Moved context generation to base agent instead of system_metrics

sia/base_agent.py

@@ -1,6 +1,7 @@
 from abc import ABC, abstractmethod
 from typing import List
 import xml.etree.ElementTree as ET
+import time
 
 from .llm_engine import LlmEngine
 from .response_parser import ResponseParser
@@ -27,7 +28,6 @@ class BaseAgent(ABC):
         """
         Initialize agent with required components.
         """
-        # Initialize components
         self._working_memory = working_memory
         self._metrics = system_metrics
         self._llm = llm
@@ -48,9 +48,27 @@ class BaseAgent(ABC):
         Returns:
             str: Complete context as XML string
         """
+        # Get memory context and calculate size
         memory_context = self._working_memory.generate_context()
         context_size = len(memory_context) / 100
-        context = self._metrics.generate_context(context_size, self._parser.io_buffer.buffer_length())
+        
+        # Get system metrics
+        metrics_data = self._metrics.get_metrics()
+        
+        # Create context element with metrics
+        context = ET.Element("context")
+        context.set("time", metrics_data["timestamp"])
+        context.set("cpu", str(metrics_data["cpu"]))
+        context.set("gpu", str(metrics_data["gpu"]))
+        context.set("memory_used", str(metrics_data["memory_used"]))
+        context.set("memory_total", str(metrics_data["memory_total"]))
+        context.set("disk_used", str(metrics_data["disk_used"]))
+        context.set("disk_total", str(metrics_data["disk_total"]))
+        context.set("context", str(round(context_size * 100)))
+        context.set("stdin", str(self._parser.io_buffer.buffer_length()))
+        
+        # Add memory entries
         for entry in memory_context:
             context.append(entry)
-        return pretty_print_element(context)
+            
+        return pretty_print_element(context)

sia/system_metrics.py

@@ -1,19 +1,18 @@
 import threading
 import time
-from typing import Optional
-import xml.etree.ElementTree as ET
+from typing import Optional, Dict
 import psutil
 
 class SystemMetrics:
     """
     Tracks system metrics including CPU, GPU, memory and disk usage.
-    Maintains rolling averages for CPU and GPU between context generations.
+    Maintains rolling averages for CPU and GPU.
     
     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
+        _cpu_samples: List of CPU usage samples since last reading
+        _gpu_samples: List of GPU usage samples since last reading
         _metrics_lock: Lock for thread-safe access to samples
     """
     
@@ -69,52 +68,47 @@ class SystemMetrics:
                     
             time.sleep(self._sample_interval)
 
-    def generate_context(self, context_usage: float, stdin_buffer_length: int) -> ET.Element:
+    def get_metrics(self) -> Dict:
         """
-        Generate XML context element with current system metrics.
-        Clears usage samples after generating averages.
+        Get current system metrics.
+        Clears usage samples after calculating averages.
         
-        Args:
-            context_usage: Current context size usage (0-1)
-            
         Returns:
-            ET.Element: Context element with system metrics
+            Dict containing:
+                - cpu: Average CPU usage percentage
+                - gpu: Average GPU usage percentage
+                - memory_used: Used memory in bytes
+                - memory_total: Total memory in bytes
+                - disk_used: Used disk space in bytes
+                - disk_total: Total disk space in bytes
+                - timestamp: Current UTC timestamp in ISO format
         """
-        # Create context element
-        context = ET.Element("context")
+        metrics = {}
         
         # Add timestamp
-        context.set("time", time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()))
+        metrics["timestamp"] = 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))
+            metrics["cpu"] = round(sum(self._cpu_samples) / len(self._cpu_samples)) if self._cpu_samples else 0
             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))
+            metrics["gpu"] = round(sum(self._gpu_samples) / len(self._gpu_samples)) if self._gpu_samples else 0
             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))
+        metrics["memory_used"] = memory.used
+        metrics["memory_total"] = 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))
+        metrics["disk_used"] = disk.used
+        metrics["disk_total"] = disk.total
         
-        # Context usage (0-100)
-        context.set("context_usage", str(round(context_usage * 100)))
-        
-        # Standard input buffer size 
-        context.set("stdin", stdin_buffer_length)
-        
-        return context
+        return metrics
         
     def stop(self):
         """Stop the monitoring thread and clean up."""