47 lines
1.6 KiB
Python
47 lines
1.6 KiB
Python
from .inference_result import InferenceResult
|
|
|
|
import xml.etree.ElementTree as ET
|
|
import re
|
|
|
|
def get_valid_root_elements(schema: str) -> set:
|
|
"""
|
|
Extract valid root element names from the XML schema.
|
|
|
|
Args:
|
|
schema: XML schema string
|
|
|
|
Returns:
|
|
set: Set of valid root element names
|
|
"""
|
|
try:
|
|
schema = schema.strip()
|
|
ET.register_namespace('xs', 'http://www.w3.org/2001/XMLSchema')
|
|
root = ET.fromstring(schema)
|
|
ns = {'xs': 'http://www.w3.org/2001/XMLSchema'}
|
|
elements = root.findall(".//xs:element", ns)
|
|
return {elem.get('name') for elem in elements if elem.get('name')}
|
|
except ET.ParseError as e:
|
|
print(f"Error parsing schema: {e}")
|
|
return set()
|
|
|
|
def split_response(response: str, valid_elements: set) -> InferenceResult:
|
|
"""
|
|
Split the response into reasoning and actions based on valid XML elements.
|
|
|
|
Args:
|
|
response: Raw response string from the model
|
|
valid_elements: Set of valid root element names from the schema
|
|
|
|
Returns:
|
|
InferenceResult: Tuple containing reasoning and actions
|
|
"""
|
|
elements_pattern = '|'.join(map(re.escape, valid_elements))
|
|
pattern = f"<({elements_pattern})[^>]*>"
|
|
matches = list(re.finditer(pattern, response))
|
|
if not matches:
|
|
return InferenceResult(response.strip(), "")
|
|
last_match = matches[-1]
|
|
split_point = last_match.start()
|
|
reasoning = response[:split_point].strip()
|
|
actions = response[split_point:].strip()
|
|
return InferenceResult(reasoning, actions) |