Tool families: a rendered family parameter keeps its identity

graphcore/tools/schemas.py

@@ -188,9 +188,19 @@ def map_type[T, U](t: Any, to_rewrite: type[T], f: Callable[[type[T]], type[U]])
     return origin[new_args]
 
 class _TemplatedTool[T: type[BaseModel], M: ToolFamilyParams, **P](BaseModel):
+    """A schema whose prose carries `{placeholder}`s, paired with the params that name them.
+
+    :meth:`with_template` renders it into the schema an LLM is actually shown. The two variants
+    below differ only in what that rendered class derives from."""
+
     _wrapped: ClassVar[type[BaseModel]]
     _key_type: ClassVar[type[ToolFamilyParams]]
 
+    @classmethod
+    def _render_onto(cls) -> type[BaseModel]:
+        """The base a rendered schema derives from."""
+        raise NotImplementedError
+
     @classmethod
     def with_template(cls, *args: P.args, **kwargs: P.kwargs) -> T:
         assert cls._wrapped.__doc__ is not None
@@ -215,33 +225,65 @@ def type_mapper(
                 new_fields[k] = (Annotated[actual_type, *descr["metadata"]], Field(**new_attrs))
             else:
                 new_fields[k] = (actual_type, Field(**new_attrs))
+        onto = cls._render_onto()
         return create_model(
             cls._wrapped.__name__,
             __doc__=new_doc,
-            __base__=cast(T, cls._wrapped),
+            __base__=cast(T, onto),
+            # JsonPlus restores by importing module+name; create_model would
+            # otherwise claim this file, which no such name admits.
+            __module__=onto.__module__,
             **new_fields
         )
 
+
+class _ToolFamily[T: type[BaseModel], M: ToolFamilyParams, **P](_TemplatedTool[T, M, P]):
+    """The handle :func:`tool_family` binds: a stand-in for the family, not a schema of its own.
+
+    It is never a value's type, so a rendering is just the wrapped schema."""
+
+    @override
+    @classmethod
+    def _render_onto(cls) -> type[BaseModel]:
+        return cls._wrapped
+
     @staticmethod
-    def _for_type[X: BaseModel, K: ToolFamilyParams,  **R](t: type[X], m: type[K]) -> type["_TemplatedTool[type[X], K, R]"]:
+    def of[X: BaseModel, K: ToolFamilyParams,  **R](t: type[X], m: type[K]) -> type["_ToolFamily[type[X], K, R]"]:
         clone = create_model(
             f"{t.__name__}Template",
-            __base__=(_TemplatedTool,)
+            __base__=(_ToolFamily,),
+            __module__=t.__module__
         )
         clone._wrapped = t
         clone._key_type = m
 
         return clone
 
+
+class _FamilyParam[T: type[BaseModel], M: ToolFamilyParams, **P](_TemplatedTool[T, M, P]):
+    """The class :func:`family_param` binds: a subclass of the wrapped schema, so it is a usable
+    annotation, and a rendering of it is a subclass of *this*.
+
+    A rendering claims this class's ``__module__`` and ``__name__``. LangGraph restores a
+    model by importing its class, which only this (module-level) name admits."""
+
+    @override
+    @classmethod
+    def _render_onto(cls) -> type[BaseModel]:
+        return cls
+
     @staticmethod
-    def _for_param_type[X: BaseModel, K: ToolFamilyParams,  **R](t: type[X], m: type[K]):
+    def of[X: BaseModel, K: ToolFamilyParams,  **R](t: type[X], m: type[K]) -> type[X]:
         clone = create_model(
             t.__name__,
-            __base__=(t, _TemplatedTool),
-            __doc__=t.__doc__
+            __base__=(t, _FamilyParam),
+            __doc__=t.__doc__,
+            # The decorator binds this class to `t`'s name in `t`'s module, so that is where it
+            # lives; create_model would otherwise have it claim this one.
+            __module__=t.__module__
         )
 
-        clone_narrowed = cast(type[_TemplatedTool[type[X], K, R]], clone)
+        clone_narrowed = cast(type[_FamilyParam[type[X], K, R]], clone)
         clone_narrowed._wrapped = t
         clone_narrowed._key_type = m
 
@@ -292,7 +334,7 @@ def family_param[
     m: Callable[P, M]
 ) -> Callable[[type[T]], type[T]]:
     def wrapper(t: type[T]):
-        return _map_templated_type(m, t, _TemplatedTool._for_param_type)
+        return _map_templated_type(m, t, _FamilyParam.of)
     return wrapper
 
 def tool_family[
@@ -303,5 +345,5 @@ def tool_family[
     m: Callable[P, M],
 ) -> Callable[[type[T]], type[_TemplatedTool[type[T], M, P]]]:
     def wrapper(t: type[T]):
-        return cast(type[_TemplatedTool[type[T], M, P]], _map_templated_type(m, t, _TemplatedTool._for_type))
+        return cast(type[_TemplatedTool[type[T], M, P]], _map_templated_type(m, t, _ToolFamily.of))
     return wrapper

tests/test_tool_families.py

@@ -1,4 +1,6 @@
 # pyright: reportInvalidTypeForm=false
+import importlib
+
 import pytest
 
 from pydantic import BaseModel, Field, create_model, ValidationError
@@ -8,8 +10,10 @@
 
 from hypothesis import HealthCheck, given, settings, strategies as st, Phase
 
+from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
+
 from graphcore.tools.schemas import (
-    WithImplementation, WithInjectedState, ToolFamilyParams, tool_family,
+    WithImplementation, WithInjectedState, ToolFamilyParams, family_param, tool_family,
 )
 
 class TemplateArgValues(TypedDict):
@@ -291,6 +295,116 @@ class Pantry(WithImplementation):
     assert staple_ty.__doc__ == "An ingredient of the curry"
 
 
+# ---------------------------------------------------------------------------
+# `family_param`: the class the decorator binds is a usable annotation, and it keeps its
+# identity -- a rendered instance is an instance of it, and it lives where it was declared.
+# ---------------------------------------------------------------------------
+
+@family_param(RecipeParams)
+class Portion(BaseModel):
+    """A portion of the {dish}"""
+    grams: int = Field(description="How many grams of the {dish} to serve")
+
+
+class ServeDish(WithImplementation):
+    """Serve the {dish}"""
+    portion: Portion = Field(description="The portion of {dish} to plate")
+
+
+class Meal(BaseModel):
+    """Where a value the tool built is stored afterwards, annotated with the bound name."""
+    portions: list[Portion]
+
+
+def test_family_param_renders_a_subtype_of_the_bound_name():
+    served = tool_family(RecipeParams)(ServeDish).with_template(dish="risotto")
+
+    portion_ty = served.model_fields["portion"].annotation
+    assert isinstance(portion_ty, type)
+    assert portion_ty.__doc__ == "A portion of the risotto"
+    assert portion_ty.model_fields["grams"].description == "How many grams of the risotto to serve"
+    assert issubclass(portion_ty, Portion)
+    assert portion_ty.__name__ == "Portion"
+    assert portion_ty.__module__ == Portion.__module__
+
+
+def test_family_param_value_validates_against_the_bound_name():
+    served = tool_family(RecipeParams)(ServeDish).with_template(dish="stew")
+
+    plated = served.model_validate({"portion": {"grams": 200}})
+    assert isinstance(plated.portion, Portion)
+
+    stored = Meal.model_validate({"portions": [plated.portion]})
+    assert stored.portions[0].grams == 200
+    assert isinstance(stored.portions[0], Portion)
+
+
+def test_family_param_lives_where_the_decorator_bound_it():
+    # What a checkpoint serializer needs: it restores a model by importing its class.
+    module = importlib.import_module(Portion.__module__)
+    assert getattr(module, Portion.__name__) is Portion
+
+    rendered = Portion.with_template(dish="curry")  # type: ignore[attributeAccessIssue]
+    assert rendered.__module__ == Portion.__module__
+    assert rendered.__name__ == Portion.__name__
+    # The rendering claims that location; lookup still returns the bound class.
+    assert getattr(module, rendered.__name__) is Portion
+
+
+def test_a_rendered_value_survives_a_checkpoint_round_trip():
+    # JsonPlus names a model by module and class. A rendering claims the bound
+    # class's location, so restore constructs that class, not a dict.
+    class Plate(WithImplementation):
+        """Plate the {dish}"""
+        portion: Portion = Field(description="The portion of {dish} to plate")
+        def run(self) -> Portion:
+            return self.portion
+
+    tool = tool_family(RecipeParams)(Plate).with_template(dish="risotto").as_tool("plate")
+    plated = tool.invoke({"portion": {"grams": 200}})
+    assert isinstance(plated, Portion)
+    assert type(plated) is not Portion
+
+    serde = JsonPlusSerializer()
+    (restored,) = serde.loads_typed(serde.dumps_typed([plated]))
+    assert type(restored) is Portion, f"restored as {type(restored)}, not the bound class"
+    assert restored.grams == 200
+
+    restored_list = serde.loads_typed(serde.dumps_typed([plated, plated]))
+    assert [type(x) is Portion and x.grams == 200 for x in restored_list] == [True, True]
+
+    rendered = tool_family(RecipeParams)(ServeDish).with_template(dish="stew")
+    raw = rendered.model_validate({"portion": {"grams": 50}}).portion
+    assert type(raw) is not Portion
+    (restored_raw,) = serde.loads_typed(serde.dumps_typed([raw]))
+    assert type(restored_raw) is Portion
+    assert restored_raw.grams == 50
+
+
+def test_family_param_renderings_stay_independent():
+    stew = tool_family(RecipeParams)(ServeDish).with_template(dish="stew")
+    pie = tool_family(RecipeParams)(ServeDish).with_template(dish="pie")
+
+    stew_portion = stew.model_fields["portion"].annotation
+    pie_portion = pie.model_fields["portion"].annotation
+    assert isinstance(stew_portion, type) and isinstance(pie_portion, type)
+    assert stew_portion is not pie_portion
+    assert issubclass(stew_portion, Portion) and issubclass(pie_portion, Portion)
+    assert not issubclass(stew_portion, pie_portion)
+    assert stew_portion.__doc__ == "A portion of the stew"
+    assert pie_portion.__doc__ == "A portion of the pie"
+
+
+def test_family_param_is_directly_renderable():
+    portion = Portion.with_template(dish="curry")  # type: ignore[attributeAccessIssue]
+
+    assert issubclass(portion, Portion)
+    assert portion.__doc__ == "A portion of the curry"
+    serde = JsonPlusSerializer()
+    restored = serde.loads_typed(serde.dumps_typed(portion(grams=3)))
+    assert type(restored) is Portion
+
+
 def test_inconsistent_key_types_rejected():
     class GardenParams(ToolFamilyParams):
         dish: str