0.1.7 - Improve nested models and attr/element support

README.md

@@ -68,6 +68,32 @@ print(f"{chat.last!r}")
 # Message(role='assistant', parts=[], content='new content')
 ```
 
+### Complex Models
+
+```python
+import rigging as rg
+
+class Inner(rg.Model):
+    type: str = rg.attr()
+    content: str
+
+class Outer(rg.Model):
+    name: str = rg.attr()
+    inners: list[Inner] = rg.element()
+
+outer = Outer(name="foo", inners=[
+    Inner(type="cat", content="meow"),
+    Inner(type="dog", content="bark")
+])
+
+print(outer.to_pretty_xml())
+
+# <outer name="foo">
+#    <inner type="cat">meow</inner>
+#    <inner type="dog">bark</inner>
+# </outer>
+```
+
 ### Tools
 
 ```python

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "rigging"
-version = "0.1.6"
+version = "0.1.7"
 description = "LLM Interaction Framework"
 authors = ["Nick Landers <monoxgas@gmail.com>"]
 license = "MIT"

rigging/__init__.py

@@ -1,9 +1,9 @@
 from rigging.generator import get_generator
 from rigging.message import Message, MessageDict, Messages
-from rigging.model import Model
+from rigging.model import Model, attr, element
 from rigging.tool import Tool
 
-__all__ = ["get_generator", "Message", "MessageDict", "Messages", "Tool", "Model"]
+__all__ = ["get_generator", "Message", "MessageDict", "Messages", "Tool", "Model", "attr", "element"]
 
 from loguru import logger
 

rigging/model.py

@@ -5,6 +5,8 @@
 from pydantic import ValidationError, field_validator
 from pydantic.alias_generators import to_snake
 from pydantic_xml import BaseXmlModel
+from pydantic_xml import attr as attr
+from pydantic_xml import element as element
 from pydantic_xml.element import SearchMode  # type: ignore [attr-defined]
 from pydantic_xml.typedefs import NsMap
 
@@ -21,6 +23,8 @@
 # content to be escaped. We should probably just write something
 # custom for our use case that supports JSON, YAML, and XML
 
+BASIC_TYPES = [int, str, float, bool]
+
 
 class XmlTagDescriptor:
     def __get__(self, _: t.Any, owner: t.Any) -> str:
@@ -68,7 +72,7 @@ def to_pretty_xml(self) -> str:
     @classmethod
     def is_simple(cls) -> bool:
         field_values = list(cls.model_fields.values())
-        return len(field_values) == 1
+        return len(field_values) == 1 and field_values[0].annotation in BASIC_TYPES
 
     @classmethod
     def xml_start_tag(cls) -> str:
@@ -100,7 +104,7 @@ def xml_example(cls) -> str:
 
     @classmethod
     def extract_xml(cls, content: str) -> tuple[ModelGeneric, str]:
-        pattern = r"(<([\w-]+)>((.*?)</\2>))"
+        pattern = r"(<([\w-]+).*?>((.*?)</\2>))"
 
         matches = re.findall(pattern, content, flags=re.DOTALL)
         matches_with_tag = [m for m in matches if m[1] == cls.__xml_tag__]

tests/test_xml_parsing.py

@@ -3,7 +3,31 @@
 
 import pytest
 
-from rigging.model import Answer, CommaDelimitedAnswer, DelimitedAnswer, Model, Question, QuestionAnswer, YesNoAnswer
+from rigging.model import (
+    Answer,
+    CommaDelimitedAnswer,
+    DelimitedAnswer,
+    Model,
+    Question,
+    QuestionAnswer,
+    YesNoAnswer,
+    attr,
+    element,
+)
+
+
+class NameWithThings(Model):
+    name: str = attr()
+    things: list[str] = element("thing")
+
+
+class Inner(Model):
+    type: str = attr()
+    content: str
+
+
+class Wrapped(Model):
+    inners: list[Inner] = element()
 
 
 @pytest.mark.parametrize(
@@ -61,6 +85,14 @@
             "<delimited-answer>hello / world / foo / bar, test | value</delimited-answer>",
             [(DelimitedAnswer, ["hello", "world", "foo", "bar, test | value"])],
         ),
+        pytest.param(
+            '<name-with-things name="test"><thing>a</thing><thing>b</thing></name-with-things>',
+            [(NameWithThings, NameWithThings(name="test", things=["a", "b"]))],
+        ),
+        pytest.param(
+            '<wrapped><inner type="cat">meow</inner><inner type="dog">bark</inner></wrapped>',
+            [(Wrapped, Wrapped(inners=[Inner(type="cat", content="meow"), Inner(type="dog", content="bark")]))],
+        ),
     ],
 )
 def test_xml_parsing(content: str, expectations: list[tuple[Model, str]]) -> None: