diff --git a/pydantic_ai_slim/pydantic_ai/prompt.py b/pydantic_ai_slim/pydantic_ai/prompt.py index 67c849d7..6dc2d336 100644 --- a/pydantic_ai_slim/pydantic_ai/prompt.py +++ b/pydantic_ai_slim/pydantic_ai/prompt.py @@ -189,29 +189,46 @@ def __init__(self, tag: str, content: Any): self._tag = tag self._content = prepare_content(content) - def build(self) -> str: + def build(self, indent: bool = True) -> str: """Build the XML string representation of the content. + Args: + indent: If True, format the output with 2-space indentation. + Default is True. + Returns: A string containing the XML representation of the content. Special characters are properly escaped, and boolean values are converted to lowercase 'true'/'false'. Examples: - >>> XMLTagBuilder('user', {'name': 'John'}).build() - 'John' + >>> print(XMLTagBuilder('user', {'name': 'John'}).build(indent=False)) + John + + >>> print(XMLTagBuilder('user', {'name': 'John'}).build()) + + John + - >>> XMLTagBuilder('items', [1, 2]).build() - '12' + >>> print(XMLTagBuilder('items', [1, 2]).build()) + 1 + 2 """ elements = self._build_element(self._tag, self._content) if isinstance(elements, list): + if indent: + return '\n'.join( + ET.indent(element, space=' ') or ET.tostring(element, encoding='unicode', method='xml').strip() + for element in cast(list[ET.Element], elements) + ) return ''.join( ET.tostring(element, encoding='unicode', method='xml').strip() for element in cast(list[ET.Element], elements) ) + if indent: + ET.indent(elements, space=' ') return ET.tostring(elements, encoding='unicode', method='xml').strip() def _build_element(self, tag: str, content: Content) -> XMLContent: diff --git a/tests/test_prompt.py b/tests/test_prompt.py index d81c7940..8ccc004d 100644 --- a/tests/test_prompt.py +++ b/tests/test_prompt.py @@ -233,7 +233,7 @@ class TestXMLContentFormatting: ) def test_simple_schema(self, content: Any) -> None: builder = XMLTagBuilder('context', content) - got = builder.build() + got = builder.build(indent=False) assert got == '42UKLondon' @@ -256,7 +256,7 @@ def test_simple_schema(self, content: Any) -> None: ) def test_list(self, content: Any) -> None: builder = XMLTagBuilder('examples', content) - got = builder.build() + got = builder.build(indent=False) assert ( got @@ -268,23 +268,54 @@ def test_dict_with_list(self) -> None: 'context', {'users': ['John', 'Jane']}, ) - got = builder.build() + got = builder.build(indent=False) assert got == 'JohnJane' def test_str(self) -> None: rules = ['Rule #1', 'Rule #2'] builder = XMLTagBuilder('rules', '\n'.join(f'- {rule}' for rule in rules)) - got = builder.build() + got = builder.build(indent=False) assert got == '- Rule #1\n- Rule #2' def test_escaping(self) -> None: builder = XMLTagBuilder('user', {'name': 'John & Jane', 'age': 30}) - got = builder.build() + got = builder.build(indent=False) assert got == '</name>John & Jane30' + def test_indent(self) -> None: + builder = XMLTagBuilder('user', {'name': 'John', 'age': 30}) + got = builder.build(indent=True) + + assert got == '\n'.join( + line + for line in ( + '', + ' John', + ' 30', + '', + ) + ) + + def test_indent_list(self) -> None: + builder = XMLTagBuilder( + 'context', + {'users': ['John', 'Jane']}, + ) + got = builder.build(indent=True) + + assert got == '\n'.join( + line + for line in ( + '', + ' John', + ' Jane', + '', + ) + ) + @pytest.mark.parametrize( ['tag', 'function'],