135 lines
5.7 KiB
Python
135 lines
5.7 KiB
Python
from datetime import datetime
|
|
from pathlib import Path
|
|
import xml.etree.ElementTree as ET
|
|
from typing import Union
|
|
|
|
from .util import format_timestamp
|
|
|
|
from .command import Command
|
|
from .delete_command import DeleteCommand
|
|
from .stop_command import StopCommand
|
|
|
|
from .entry.background_entry import BackgroundEntry
|
|
from .entry import Entry
|
|
from .io_buffer import IOBuffer
|
|
from .entry.parse_error_entry import ParseErrorEntry
|
|
from .entry.read_entry import ReadEntry
|
|
from .entry.reasoning_entry import ReasoningEntry
|
|
from .entry.repeat_entry import RepeatEntry
|
|
from .entry.single_entry import SingleEntry
|
|
from .entry.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, work_dir: Path, io_buffer: IOBuffer):
|
|
"""
|
|
Initialize parser with IO buffer.
|
|
|
|
Args:
|
|
work_dir: Workdir for the scripts
|
|
io_buffer: Buffer to use for IO operations
|
|
"""
|
|
self._work_dir = work_dir
|
|
self._io_buffer = io_buffer
|
|
|
|
@property
|
|
def io_buffer(self) -> IOBuffer:
|
|
"""
|
|
Get the IO buffer used by the parser.
|
|
|
|
Returns:
|
|
IOBuffer: Buffer used for IO operations
|
|
"""
|
|
return self._io_buffer
|
|
|
|
def parse(self, timestamp: datetime, 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 = format_timestamp(timestamp)
|
|
|
|
parser = ET.XMLPullParser(events=("start", "end"))
|
|
parser.feed(xml)
|
|
root = None
|
|
try:
|
|
for event, root in parser.read_events():
|
|
if event == "start":
|
|
break
|
|
except ET.ParseError as e:
|
|
return ParseErrorEntry(entry_id, xml, f"Invalid XML: {str(e)}")
|
|
|
|
try:
|
|
if root.tag == 'delete':
|
|
target_id = root.get('id')
|
|
if not target_id:
|
|
return ParseErrorEntry(entry_id, xml, "Delete command missing required 'id' attribute")
|
|
if len(root.attrib) > 1:
|
|
return ParseErrorEntry(entry_id, xml, "Delete command should only have 'id' attribute")
|
|
return DeleteCommand(target_id)
|
|
|
|
elif root.tag == 'stop':
|
|
return StopCommand()
|
|
|
|
elif root.tag == 'background':
|
|
if len(root) != 0 or root.attrib or root.text is None or root.text.strip() == '':
|
|
return ParseErrorEntry(entry_id, xml, "Background entry requires (only) script content")
|
|
return BackgroundEntry(entry_id, self._work_dir, root.text)
|
|
|
|
elif root.tag == 'repeat':
|
|
if len(root) != 0 or root.text is None or root.text.strip() == '':
|
|
return ParseErrorEntry(entry_id, xml, "Repeat entry requires (only) script content")
|
|
if len(root.attrib) > 2 or any(k not in ('timeout', 'limit') for k in root.attrib):
|
|
return ParseErrorEntry(entry_id, xml, "Repeat entry only accepts 'timeout' and 'limit' attributes")
|
|
timeout = root.get('timeout')
|
|
limit = root.get('limit')
|
|
timeout = float(timeout) if timeout is not None else None
|
|
limit = int(limit) if limit is not None else None
|
|
return RepeatEntry(entry_id, self._work_dir, root.text, timeout, limit)
|
|
|
|
elif root.tag == 'single':
|
|
if len(root) != 0 or root.text is None or root.text.strip() == '':
|
|
return ParseErrorEntry(entry_id, xml, "Single entry requires (only) script content")
|
|
if len(root.attrib) > 2 or any(k not in ('timeout', 'limit') for k in root.attrib):
|
|
return ParseErrorEntry(entry_id, xml, "Single entry only accepts 'timeout' and 'limit' attributes")
|
|
timeout = root.get('timeout')
|
|
limit = root.get('limit')
|
|
timeout = float(timeout) if timeout is not None else None
|
|
limit = int(limit) if limit is not None else None
|
|
return SingleEntry(entry_id, self._work_dir, root.text, timeout, limit)
|
|
|
|
elif root.tag == 'reasoning':
|
|
if len(root) != 0 or root.attrib or root.text is None or root.text.strip() == '':
|
|
return ParseErrorEntry(entry_id, xml, "Reasoning entry requires (only) text content")
|
|
return ReasoningEntry(entry_id, root.text)
|
|
|
|
elif root.tag == 'read_stdin':
|
|
return ReadEntry(entry_id, self._io_buffer)
|
|
|
|
elif root.tag == 'write_stdout':
|
|
if len(root) != 0 or root.attrib or root.text is None or root.text.strip() == '':
|
|
return ParseErrorEntry(entry_id, xml, "Write stdout entry requires (only) text content")
|
|
return WriteEntry(entry_id, root.text, self._io_buffer)
|
|
|
|
else:
|
|
return ParseErrorEntry(entry_id, xml, f"Unknown root element: {root.tag}")
|
|
|
|
except Exception as e:
|
|
return ParseErrorEntry(entry_id, xml, f"Error parsing response: {str(e)}")
|