Added JSON deref for APITools. Version bump to 2.2.1.

Author: monoxgas

poetry.lock

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "rigging"
-version = "2.2.0"
+version = "2.2.1"
 description = "LLM Interaction Framework"
 authors = ["Nick Landers <monoxgas@gmail.com>"]
 license = "MIT"
@@ -22,6 +22,9 @@ colorama = "^0.4.6"
 boto3 = "^1.35.0"
 boto3-stubs = { extras = ["s3"], version = "^1.35.0" }
 logfire-api = "^3.1.1"
+jsonpath-ng = "^1.7.0"
+ruamel-yaml = "^0.18.10"
+jsonref = "^1.1.0"
 
 vllm = { version = "^0.5.0", optional = true }
 transformers = { version = "^4.41.0", optional = true }
@@ -32,8 +35,6 @@ click = { version = "^8.1.7", optional = true }
 httpx = { version = "^0.27.0", optional = true }
 aiodocker = { version = "^0.22.2", optional = true }
 websockets = { version = "^13.0", optional = true }
-jsonpath-ng = "^1.7.0"
-ruamel-yaml = "^0.18.10"
 
 [tool.poetry.extras]
 tracing = ["logfire"]

pyproject.toml

@@ -18,7 +18,7 @@
 from rigging.tool import ApiTool, Tool
 from rigging.util import await_
 
-__version__ = "2.2.0"
+__version__ = "2.2.1"
 
 __all__ = [
     "get_generator",

rigging/__init__.py

@@ -21,6 +21,7 @@
 from pydantic_xml import attr as attr
 from pydantic_xml import element as element
 from pydantic_xml import wrapped as wrapped
+from pydantic_xml.model import XmlEntityInfo
 from pydantic_xml.typedefs import EntityLocation, NsMap
 
 from rigging.error import MissingModelError
@@ -131,6 +132,11 @@ def is_simple(cls) -> bool:
         if len(field_values) != 1:
             return False
 
+        # If the field is a pydantic-xml wrapper like element()
+        # or attr(), we'll assume the user knows what they're doing
+        if isinstance(field_values[0], XmlEntityInfo):
+            return False
+
         annotation = field_values[0].annotation
         if t.get_origin(annotation) == t.Annotated:
             annotation = t.get_args(annotation)[0]

rigging/model.py

@@ -13,6 +13,7 @@
 from pydantic import BaseModel, TypeAdapter, field_validator
 
 from rigging.tracing import tracer
+from rigging.util import deref_json
 
 if t.TYPE_CHECKING:
     from rigging.message import Message
@@ -145,7 +146,7 @@ def schema(self) -> dict[str, t.Any]:
 
             schema["properties"][name]["description"] = annotation_args[1]
 
-        return schema
+        return deref_json(schema, is_json_schema=True)
 
     @cached_property
     def definition(self) -> ToolDefinition:

rigging/tool/api.py

@@ -8,6 +8,7 @@
 import typing as t
 from threading import Thread
 
+import jsonref  # type: ignore [import-untyped]
 from pydantic import alias_generators
 
 R = t.TypeVar("R")
@@ -66,6 +67,18 @@ def await_(*coros: t.Coroutine[t.Any, t.Any, R]) -> R | list[R]:  # type: ignore
     return results
 
 
+# JSON
+
+
+def deref_json(obj: dict[str, t.Any], *, is_json_schema: bool = False) -> dict[str, t.Any]:
+    return jsonref.replace_refs(  # type: ignore [no-any-return]
+        obj,
+        jsonschema=is_json_schema,
+        proxies=False,
+        lazy_load=False,
+    )
+
+
 # XML Formatting