Response parser and split up class diagrams

This commit is contained in:
2024-11-01 11:38:28 +01:00
parent 165279ecc7
commit d80171de15
6 changed files with 932 additions and 177 deletions

154
sia/response_parser.py Normal file
View File

@@ -0,0 +1,154 @@
from datetime import datetime
import uuid
import xml.etree.ElementTree as ET
from typing import Union
from .command import Command
from .delete_command import DeleteCommand
from .stop_command import StopCommand
from .entry import Entry
from .background_entry import BackgroundEntry
from .parse_error_entry import ParseErrorEntry
from .read_entry import ReadEntry
from .reasoning_entry import ReasoningEntry
from .repeat_entry import RepeatEntry
from .single_shot_entry import SingleShotEntry
from .write_entry import WriteEntry
class ResponseParser:
"""
Parses XML responses from the LLM into commands or entries.
The parser validates the XML structure and converts it into the appropriate
Command or Entry object based on the root element tag. Invalid input results
in ParseErrorEntry objects rather than raising exceptions.
Attributes:
io_buffer: Buffer to use for IO operations
"""
def __init__(self, io_buffer):
"""
Initialize parser with IO buffer.
Args:
io_buffer: Buffer to use for IO operations
"""
self.io_buffer = io_buffer
def _extract_content(self, xml: str, root: ET.Element) -> str:
"""
Extract content from XML, handling CDATA sections and preserving special characters.
Args:
xml: Original XML string
root: Parsed XML element
Returns:
str: Extracted content with proper handling of CDATA and special characters
"""
if root.text is None:
return ""
# Find opening tag end
tag_name = root.tag
tag_start = xml.find('<' + tag_name)
tag_end = xml.find('>', tag_start)
if tag_end == -1:
return ""
# Check for self-closing tag
if xml[tag_end-1] == '/':
return ""
# Find closing tag start
close_tag = '</' + tag_name + '>'
close_start = xml.rfind(close_tag)
if close_start == -1:
return ""
# Extract content
content = xml[tag_end+1:close_start]
# Handle CDATA sections
if content.startswith('<![CDATA[') and content.endswith(']]>'):
content = content[9:-3] # Remove CDATA markers
return content
def parse(self, xml: str) -> Union[Command, Entry]:
"""
Parse XML response into a Command or Entry.
Args:
xml: XML string to parse
Returns:
Command or Entry based on the XML content
ParseErrorEntry for any parsing or validation errors
"""
entry_id = str(uuid.uuid4())
timestamp = datetime.now()
# First try to parse the XML
try:
root = ET.fromstring(xml.strip())
except ET.ParseError as e:
return ParseErrorEntry(xml, f"Invalid XML: {str(e)}", entry_id, timestamp)
# Extract content preserving special characters and handling CDATA
content = self._extract_content(xml, root)
# Convert to appropriate type based on root tag
try:
if root.tag == 'delete':
# Extract target ID and create delete command
target_id = root.get('id')
if not target_id:
return ParseErrorEntry(xml, "Delete command missing required 'id' attribute", entry_id, timestamp)
return DeleteCommand(target_id)
elif root.tag == 'stop':
# Create stop command
return StopCommand()
elif root.tag == 'background':
# Create background script entry
if not content:
return ParseErrorEntry(xml, "Background entry missing script content", entry_id, timestamp)
return BackgroundEntry(content, entry_id, timestamp)
elif root.tag == 'repeat':
# Create repeat script entry
if not content:
return ParseErrorEntry(xml, "Repeat entry missing script content", entry_id, timestamp)
return RepeatEntry(content, entry_id, timestamp)
elif root.tag == 'single_shot':
# Create single shot script entry
if not content:
return ParseErrorEntry(xml, "Single shot entry missing script content", entry_id, timestamp)
return SingleShotEntry(content, entry_id, timestamp)
elif root.tag == 'reasoning':
# Create reasoning entry
if not content:
return ParseErrorEntry(xml, "Reasoning entry missing content", entry_id, timestamp)
return ReasoningEntry(content, entry_id, timestamp)
elif root.tag == 'read_stdin':
# Create read entry - no content required
return ReadEntry(self.io_buffer, entry_id, timestamp)
elif root.tag == 'write_stdout':
# Create write entry
if not content:
return ParseErrorEntry(xml, "Write entry missing content", entry_id, timestamp)
return WriteEntry(content, self.io_buffer, entry_id, timestamp)
else:
return ParseErrorEntry(xml, f"Unknown root element: {root.tag}", entry_id, timestamp)
except Exception as e:
return ParseErrorEntry(xml, f"Error parsing response: {str(e)}", entry_id, timestamp)

122
sia/system_metrics.py Normal file
View File

@@ -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()