llm/SIA
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fix pretty printing

Browse Source
This commit is contained in:
geens
2024-11-02 15:22:32 +01:00
parent a58ee40c2c
commit 96d14fc46d
6 changed files with 213 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
readme.md
View File

@@ -147,6 +147,17 @@ The LLM is free to escape data any way it wants,
as long as it results in valid XML.
The response is validated against a schema.
#### XML Data Flow
Entries store their content as raw text. During context compilation, the XML formatter
wraps text content in CDATA sections, except when the content contains CDATA closing sequences.
In those cases, the formatter uses standard XML escaping.
This separation between storage and formatting:
- Keeps entry data clean and unescaped
- Centralizes XML formatting rules
- Makes it easy to change escaping rules without modifying entries
- Allows different formatting for different use cases
The Context is escaped using CDATA blocks.
Except when the data contains CDATA closing sequences.
Then the whole block is escaped using standard XML escaping.

5
sia/base_agent.py
View File

@@ -3,10 +3,11 @@ from typing import List
import xml.etree.ElementTree as ET
from .llm_engine import LlmEngine
from .response_parser import ResponseParser
from .system_metrics import SystemMetrics
from .util import pretty_print_element
from .working_memory import WorkingMemory
from .xml_validator import XMLValidator
from .response_parser import ResponseParser
class BaseAgent(ABC):
"""
@@ -52,4 +53,4 @@ class BaseAgent(ABC):
context = self._metrics.generate_context(context_size)
for entry in memory_context:
context.append(entry)
return ET.tostring(context, encoding="unicode")
return pretty_print_element(context)

54
sia/util.py
View File

@@ -1,5 +1,6 @@
from typing import Iterator
import xml.dom.minidom
import xml.etree.ElementTree as ET
from typing import Iterator, Optional
def stop_before_value(iterator: Iterator[str], stop_value: str) -> Iterator[str]:
"""
@@ -22,6 +23,57 @@ def stop_before_value(iterator: Iterator[str], stop_value: str) -> Iterator[str]
break
yield item
def escape_text_for_xml(text: str) -> str:
"""Convert any value to string representation for XML storage."""
if text is None:
return ''
return str(text)
def pretty_print_element(elem: ET.Element, level: int = 0, max_line: int = 80) -> str:
"""Convert ElementTree element to pretty-printed string with custom formatting."""
indent = ' ' * level
# Handle empty elements
if not elem.text and not elem.tail and not elem.attrib and len(elem) == 0:
return f'{indent}<{elem.tag}/>'
# Build opening tag with attributes
tag = elem.tag
attrs = elem.attrib
if attrs:
attr_strings = [f'{k}="{v}"' for k, v in sorted(attrs.items())]
attr_str = ' ' + ' '.join(attr_strings)
if len(indent) + len(tag) + len(attr_str) + 2 > max_line:
attr_indent = indent + ' '
attr_str = '\n' + '\n'.join(f'{attr_indent}{a}' for a in attr_strings)
else:
attr_str = ''
parts = [f'{indent}<{tag}{attr_str}>']
# Handle text content
if elem.text and elem.text.strip():
text = elem.text
if ']]>' in text:
escaped = text.replace('&', '&amp;') \
.replace('<', '&lt;') \
.replace('>', '&gt;') \
.replace('"', '&quot;') \
.replace("'", '&apos;')
parts.append(f'{indent} {escaped}')
else:
parts.append(f'{indent} <![CDATA[{text}]]>')
# Handle children
for child in elem:
parts.append(pretty_print_element(child, level + 1, max_line))
if child.tail and child.tail.strip():
parts.append(f'{indent} {child.tail}')
parts.append(f'{indent}</{tag}>')
return '\n'.join(parts)
def escape_text_for_xml(text: str) -> str:
"""
Convert any text value to a properly escaped XML string.

6
test/base_agent_test.py
View File

@@ -145,7 +145,6 @@ class BaseAgentTest(unittest.TestCase):
self.assertEqual(len(list(root)), 0)
def test_compile_context_with_entries(self):
"""Test context compilation with working memory entries."""
# Add test entry to working memory
entry = ReasoningEntry("test reasoning", "test-id", self.test_timestamp)
self.agent._working_memory.add_entry(entry)
@@ -162,7 +161,8 @@ class BaseAgentTest(unittest.TestCase):
reasoning_elem = root.find("reasoning")
self.assertIsNotNone(reasoning_elem)
self.assertEqual(reasoning_elem.get("id"), "test-id")
self.assertEqual(reasoning_elem.text, "<![CDATA[test reasoning]]>")
# Strip whitespace before comparing
self.assertEqual(reasoning_elem.text.strip(), "<![CDATA[test reasoning]]>")
def test_multiple_entries(self):
"""Test handling multiple entries in working memory."""
@@ -186,4 +186,4 @@ class BaseAgentTest(unittest.TestCase):
self.assertEqual(len(reasoning_elems), 3)
for i, elem in enumerate(reasoning_elems):
self.assertEqual(elem.get("id"), f"id-{i}")
self.assertEqual(elem.text, f"<![CDATA[test {i}]]>")
self.assertEqual(elem.text.strip(), f"<![CDATA[test {i}]]>")

23
test/system_metrics_test.py
View File

@@ -138,26 +138,25 @@ class SystemMetricsTest(unittest.TestCase):
context2.get("time")
)
def test_cpu_usage_detection(self):
def test_cpu_usage_detection(self):
"""Test CPU usage increases under load"""
# Get initial CPU usage
context1 = self.metrics.generate_context(0.5)
initial_cpu = int(context1.get("cpu"))
# Get initial CPU usage - take average of multiple samples
initial_contexts = [self.metrics.generate_context(0.5) for _ in range(3)]
initial_cpu = min(int(ctx.get("cpu")) for ctx in initial_contexts)
# Generate CPU load in separate process
process = multiprocessing.Process(target=cpu_load_process)
process.start()
try:
# Wait for samples to accumulate
time.sleep(0.5)
# Wait for load to register and take multiple samples
time.sleep(1.0)
load_contexts = [self.metrics.generate_context(0.5) for _ in range(3)]
load_cpu = max(int(ctx.get("cpu")) for ctx in load_contexts)
# Get CPU usage under load
context2 = self.metrics.generate_context(0.5)
load_cpu = int(context2.get("cpu"))
# CPU usage should increase
self.assertGreater(load_cpu, initial_cpu)
# Verify load increases CPU usage
self.assertGreater(load_cpu, initial_cpu,
f"CPU usage should increase under load. Initial: {initial_cpu}, Load: {load_cpu}")
finally:
process.terminate()

121
test/util_test.py
View File

@@ -1,7 +1,6 @@
import unittest
from typing import Iterator
from . import test_data
import xml.etree.ElementTree as ET
from sia import util
@@ -116,3 +115,121 @@ class UtilTest(unittest.TestCase):
"""Test boolean values are properly converted"""
self.assertEqual(util.escape_text_for_xml(True), '<![CDATA[True]]>')
self.assertEqual(util.escape_text_for_xml(False), '<![CDATA[False]]>')
def test_pretty_print_element_basic(self):
"""Test basic element printing"""
elem = ET.Element('root')
child = ET.SubElement(elem, 'child', {'attr': 'value'})
child.text = 'text'
expected = '''<root>
<child attr="value">
<![CDATA[text]]>
</child>
</root>'''
self.assertEqual(util.pretty_print_element(elem).strip(), expected.strip())
def test_pretty_print_element_long_attributes(self):
"""Test printing elements with long attributes that need wrapping"""
elem = ET.Element('root', {
'very_long_attribute_1': 'value1',
'very_long_attribute_2': 'value2',
'very_long_attribute_3': 'value3'
})
result = util.pretty_print_element(elem, max_line=40)
# Verify attributes are wrapped to new lines
self.assertIn('\n very_long_attribute_1="value1"', result)
self.assertIn('\n very_long_attribute_2="value2"', result)
self.assertIn('\n very_long_attribute_3="value3"', result)
def test_pretty_print_element_cdata(self):
"""Test printing elements with CDATA sections"""
elem = ET.Element('root')
elem.text = 'preserved\n spacing'
expected = '''<root>
<![CDATA[preserved\n spacing]]>
</root>'''
self.assertEqual(util.pretty_print_element(elem).strip(), expected.strip())
def test_pretty_print_element_empty(self):
"""Test printing empty elements"""
elem = ET.Element('empty')
self.assertEqual(util.pretty_print_element(elem), '<empty/>')
def test_real_world_context(self):
"""Test formatting a realistic context example"""
# Create a context element with system metrics and a repeat entry
context = ET.Element('context', {
'time': '2024-10-18T12:00:00Z',
'cpu': '12',
'gpu': '26',
'memory_used': '9556302234',
'memory_total': '17179869184',
'disk_used': '244434939904',
'disk_total': '273145991168',
'context': '3',
'stdin': '0'
})
repeat = ET.SubElement(context, 'repeat', {
'id': 'a3d89ee5-28ec-4c5a-b9e9-a30af53d43a0',
'exit_code': '0'
})
repeat.text = 'ls -lah /'
stdout = ET.SubElement(repeat, 'stdout')
stdout.text = '''total 16K
drwxr-xr-x 1 sia 1049089 0 Oct 28 13:40 ./
drwxr-xr-x 1 sia 1049089 0 Oct 28 13:40 ../
drwxr-xr-x 1 sia 1049089 0 Oct 28 13:40 tasks/
drwxr-xr-x 1 sia 1049089 0 Oct 28 13:40 user/'''
stderr = ET.SubElement(repeat, 'stderr')
formatted = util.pretty_print_element(context)
# Verify key formatting expectations
self.assertIn(' <repeat exit_code="0" id="a3d89ee5-28ec-4c5a-b9e9-a30af53d43a0">', formatted)
self.assertIn(' <![CDATA[ls -lah /]]>', formatted)
self.assertIn(' <stdout>', formatted)
self.assertIn(' <![CDATA[total 16K', formatted)
def test_large_content(self):
"""Test formatting with large content blocks"""
# Create element with large text content
elem = ET.Element('root')
elem.text = 'x' * 1000
result = util.pretty_print_element(elem)
# Verify content is properly indented and wrapped
self.assertIn('x' * 1000, result)
def test_special_characters(self):
"""Test handling of special XML characters"""
elem = ET.Element('root')
child = ET.SubElement(elem, 'child')
child.text = '< > & " \''
result = util.pretty_print_element(elem)
# Should be escaped or wrapped in CDATA
self.assertTrue(
('&lt; &gt; &amp; &quot; &apos;' in result) or
('<![CDATA[< > & " \']]>' in result)
)
def test_mixed_content(self):
"""Test formatting elements with mixed text and child nodes"""
elem = ET.Element('root')
elem.text = 'before'
child = ET.SubElement(elem, 'child')
child.tail = 'after'
result = util.pretty_print_element(elem)
# Verify text positioning
self.assertIn('before', result)
self.assertIn('after', result)