69 lines
2.3 KiB
Python
69 lines
2.3 KiB
Python
from typing import Iterator, TypeVar
|
|
import re
|
|
import xml.etree.ElementTree as ET
|
|
|
|
from .inference_result import InferenceResult
|
|
|
|
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)
|
|
|
|
def stop_before_value(iterator: Iterator[str], stop_value: str) -> Iterator[str]:
|
|
"""
|
|
Creates an iterator that yields values from the input iterator
|
|
until it encounters the stop_value (exclusive).
|
|
|
|
Args:
|
|
iterator: The source iterator
|
|
stop_value: The value to stop before
|
|
|
|
Yields:
|
|
Values from the iterator until stop_value is encountered
|
|
If stop_value is part of an item, yields the part before stop_value
|
|
"""
|
|
for item in iterator:
|
|
if stop_value in item:
|
|
split_point = item.index(stop_value)
|
|
if split_point > 0:
|
|
yield item[:split_point]
|
|
break
|
|
yield item |