Transitive tool families

graphcore/tools/schemas.py

@@ -1,12 +1,16 @@
 from typing import (
-    Generic, TypeVar, Annotated, Any, ClassVar, override, Iterator, cast, Mapping, Callable, Never
+    Generic, TypeVar, Annotated, Any, ClassVar,
+    override, Iterator, cast, Callable, Never, get_args, get_origin
 )
+import types
 import typing
 import string
 import re
 from dataclasses import dataclass
 from contextlib import contextmanager
 from contextvars import ContextVar
+from functools import reduce
+from operator import or_
 
 from pydantic import BaseModel, Field, create_model
 
@@ -137,67 +141,121 @@ def tool_deps(self) -> Iterator[DEPS]:
 class InjectAll(WithInjectedState[ST], WithInjectedId):
     pass
 
-@dataclass
-class TemplatedTool[T: type[BaseModel], **P]:
-    _staged: T
 
-    def with_template(
-        self, *args: P.args, **kwargs: P.kwargs
-    ) -> T:
-        assert self._staged.__doc__ is not None
-        new_doc = self._staged.__doc__.format(*args, **kwargs)
-        assert issubclass(self._staged, BaseModel)
+@typing.dataclass_transform(kw_only_default=True)
+class ToolFamilyParams:
+    def __new__(cls) -> Never:
+        raise ValueError("These are phantom types and never meant to be instantiated")
+
+def _placeholders(fmt: str) -> set[str]:
+    to_ret = set()
+    for _, fn, _, _ in string.Formatter().parse(fmt):
+        if not fn:
+            continue
+        if re.match(r"^[A-Za-z_][A-Za-z0-9_]*$", fn) is None:
+            raise ValueError("Cannot define non-simple placeholder")
+        to_ret.add(fn)
+    return to_ret
+
+
+def map_type[T, U](t: Any, to_rewrite: type[T], f: Callable[[type[T]], type[U]]) -> Any:
+    """Rebuild type expression `t`, replacing occurrences of `to_rewrite` with f(match)."""
+    # A match rewrites and stops -- we don't descend into the matched type.
+    if isinstance(t, type) and issubclass(t, to_rewrite):
+        return f(t)
+
+    origin = get_origin(t)
+    if origin is None:
+        return t  # leaf: plain class, None, Ellipsis, a Literal value, ...
+
+    # Annotated[X, meta...]: walk X, leave metadata alone.
+    if origin is Annotated:
+        inner, *meta = get_args(t)
+        return Annotated[tuple([map_type(inner, to_rewrite, f), *meta])]
+
+    args = get_args(t)
+    new_args = tuple(
+        [map_type(x, to_rewrite, f) for x in a] if isinstance(a, list)  # Callable's [params]
+        else map_type(a, to_rewrite, f)
+        for a in args
+    )
+    if new_args == args:
+        return t  # untouched subtree: hand back the original object
+
+    if origin is types.UnionType:  # X | Y can't be rebuilt as origin[args]
+        return reduce(or_, new_args)
+
+    return origin[new_args]
+
+class _TemplatedTool[T: type[BaseModel], M: ToolFamilyParams, **P](BaseModel):
+    _wrapped: ClassVar[type[BaseModel]]
+    _key_type: ClassVar[type[ToolFamilyParams]]
+
+    @classmethod
+    def with_template(cls, *args: P.args, **kwargs: P.kwargs) -> T:
+        assert cls._wrapped.__doc__ is not None
+        new_doc = cls._wrapped.__doc__.format(*args, **kwargs)
+        assert issubclass(cls._wrapped, BaseModel)
         new_fields : dict[str, Any] = {}
-        for (k, v) in self._staged.model_fields.items():
-            if not v.description:
-                continue
+        def type_mapper(
+            t: type[_TemplatedTool]
+        ) -> type[Any]:
+            return t.with_template(*args, **kwargs)
+        for (k, v) in cls._wrapped.model_fields.items():
+            actual_type = v.annotation
+            if v.annotation is not None:
+                actual_type = map_type(v.annotation, _TemplatedTool, type_mapper)
             descr = v.asdict()
             new_attrs = {
                 **descr["attributes"],
-                "description": v.description.format(*args, **kwargs)
             }
+            if v.description:
+                new_attrs["description"] = v.description.format(*args, **kwargs)
             if descr["metadata"]:
-                new_fields[k] = (Annotated[v.annotation, *descr["metadata"]], Field(**new_attrs))
+                new_fields[k] = (Annotated[actual_type, *descr["metadata"]], Field(**new_attrs))
             else:
-                new_fields[k] = (v.annotation, Field(**new_attrs))
+                new_fields[k] = (actual_type, Field(**new_attrs))
         return create_model(
-            f"{self._staged.__name__}Templated",
+            cls._wrapped.__name__,
             __doc__=new_doc,
-            __base__=self._staged,
+            __base__=cast(T, cls._wrapped),
             **new_fields
         )
 
-def _placeholders(fmt: str) -> set[str]:
-    to_ret = set()
-    for _, fn, _, _ in string.Formatter().parse(fmt):
-        if not fn:
-            continue
-        if re.match(r"^[A-Za-z_][A-Za-z0-9_]*$", fn) is None:
-            raise ValueError("Cannot define non-simple placeholder")
-        to_ret.add(fn)
-    return to_ret
-
+    @staticmethod
+    def _for_type[X: BaseModel, K: ToolFamilyParams,  **R](t: type[X], m: type[K]) -> type["_TemplatedTool[type[X], K, R]"]:
+        clone = create_model(
+            f"{t.__name__}Template",
+            __base__=(_TemplatedTool,)
+        )
+        clone._wrapped = t
+        clone._key_type = m
 
-@typing.dataclass_transform(kw_only_default=True)
-class ToolFamilyParams:
-    def __new__(cls) -> Never:
-        raise ValueError("These are phantom types and never meant to be instantiated")
+        return clone
 
-def tool_family[T:
-    WithAsyncDependencies | WithAsyncImplementation | WithImplementation,
+def tool_family[
+    T: BaseModel,
     M: ToolFamilyParams,
     **P,
 ](
     m: Callable[P, M],
-) -> Callable[[type[T]], TemplatedTool[type[T], P]]:
+) -> Callable[[type[T]], type[_TemplatedTool[type[T], M, P]]]:
     def wrapper(t: type[T]):
         assert isinstance(m, type)
         assert issubclass(m, ToolFamilyParams) and issubclass(t, BaseModel)
         doc = t.__doc__
         assert doc is not None
         params = set()
         params |= _placeholders(doc)
+        def check_key(nested: type[_TemplatedTool]) -> type[_TemplatedTool]:
+            if nested._key_type is not m:
+                raise ValueError(
+                    f"Cannot use inconsistent key types: {m} vs {nested._key_type} (via {nested.__name__})"
+                )
+            return nested
         for (k, v) in t.model_fields.items():
+            if v.annotation is not None:
+                map_type(v.annotation, _TemplatedTool, check_key)
             if not v.description:
                 continue
             params |= _placeholders(v.description)
@@ -206,5 +264,5 @@ def wrapper(t: type[T]):
             missing = params - annots.keys()
             if missing:
                 raise ValueError(f"Missing declared tool params: {missing}")
-        return TemplatedTool(t)
+        return _TemplatedTool._for_type(t, cast(type[M], m))
     return wrapper

tests/test_tool_families.py

@@ -1,5 +1,8 @@
+# pyright: reportInvalidTypeForm=false
+import pytest
+
 from pydantic import BaseModel, Field, create_model, ValidationError
-from typing import Annotated, cast, Any, TypedDict
+from typing import Annotated, cast, Any, TypedDict, get_args, get_origin
 
 from annotated_types import Gt, Ge, Le, Lt
 
@@ -184,6 +187,133 @@ def test_tool_family_preserves_validation(data: st.DataObject) -> None:
     assert validation_outcome(basic, payload) == validation_outcome(templated, payload)
 
 
+# ---------------------------------------------------------------------------
+# Transitive templating: fields whose annotations mention another tool family
+# are templated with the same arguments as the enclosing family.
+# ---------------------------------------------------------------------------
+
+class RecipeParams(ToolFamilyParams):
+    dish: str
+
+
+class Ingredient(BaseModel):
+    """An ingredient of the {dish}"""
+    name: str = Field(description="Name of the ingredient in the {dish}")
+    amount: int = Field(description="How much of it to use")
+
+
+IngredientFamily = tool_family(RecipeParams)(Ingredient)
+
+
+class MakeRecipe(WithImplementation):
+    """Write a recipe for the {dish}"""
+    title: str = Field(description="Title of the {dish} recipe")
+    main: IngredientFamily = Field(description="The main ingredient") #type: ignore[invalidTypeForm]
+    extras: list[IngredientFamily] = Field(description="Additional ingredients")
+    garnish: IngredientFamily | None = Field(default=None, description="Optional garnish")
+
+
+RecipeFamily = tool_family(RecipeParams)(MakeRecipe)
+
+
+def test_transitive_template_direct_field():
+    recipe = RecipeFamily.with_template(dish="paella")
+
+    assert recipe.__name__ == "MakeRecipe"
+    assert recipe.__doc__ == "Write a recipe for the paella"
+    assert recipe.model_fields["title"].annotation is str
+    assert recipe.model_fields["title"].description == "Title of the paella recipe"
+
+    main_ty = recipe.model_fields["main"].annotation
+    assert isinstance(main_ty, type) and issubclass(main_ty, Ingredient)
+    assert main_ty.__doc__ == "An ingredient of the paella"
+    assert main_ty.model_fields["name"].description == "Name of the ingredient in the paella"
+
+
+def test_transitive_template_inside_containers():
+    recipe = RecipeFamily.with_template(dish="soup")
+
+    extras_ty = recipe.model_fields["extras"].annotation
+    assert get_origin(extras_ty) is list
+    (elem_ty,) = get_args(extras_ty)
+    assert issubclass(elem_ty, Ingredient)
+    assert elem_ty.__doc__ == "An ingredient of the soup"
+
+    garnish_ty = recipe.model_fields["garnish"].annotation
+    garnish_args = get_args(garnish_ty)
+    assert type(None) in garnish_args
+    (inner_ty,) = [a for a in garnish_args if a is not type(None)]
+    assert issubclass(inner_ty, Ingredient)
+    assert inner_ty.__doc__ == "An ingredient of the soup"
+
+
+def test_transitive_template_validation():
+    recipe = RecipeFamily.with_template(dish="stew")
+
+    parsed = recipe.model_validate({
+        "title": "Beef stew",
+        "main": {"name": "beef", "amount": 2},
+        "extras": [{"name": "carrot", "amount": 3}],
+        "garnish": None,
+    })
+    assert parsed.main.amount == 2
+    assert parsed.extras[0].name == "carrot"
+
+    with pytest.raises(ValidationError):
+        recipe.model_validate({
+            "title": "Beef stew",
+            "main": {"name": "beef"},  # missing amount
+            "extras": [],
+            "garnish": None,
+        })
+
+
+def test_templated_instances_are_independent():
+    soup = RecipeFamily.with_template(dish="soup")
+    pie = RecipeFamily.with_template(dish="pie")
+
+    soup_main = soup.model_fields["main"].annotation
+    pie_main = pie.model_fields["main"].annotation
+    assert soup_main is not pie_main
+    assert soup_main.__doc__ == "An ingredient of the soup"
+    assert pie_main.__doc__ == "An ingredient of the pie"
+
+
+def test_transitive_template_without_field_description():
+    class Pantry(WithImplementation):
+        """Check the pantry for the {dish}"""
+        staple: IngredientFamily
+
+    pantry = tool_family(RecipeParams)(Pantry).with_template(dish="curry")
+
+    staple_ty = pantry.model_fields["staple"].annotation
+    assert isinstance(staple_ty, type) and issubclass(staple_ty, Ingredient)
+    assert staple_ty.__doc__ == "An ingredient of the curry"
+
+
+def test_inconsistent_key_types_rejected():
+    class GardenParams(ToolFamilyParams):
+        dish: str
+
+    with pytest.raises(ValueError, match="inconsistent key types"):
+        @tool_family(GardenParams)
+        class BadRecipe(WithImplementation):
+            """Write a recipe for the {dish}"""
+            main: IngredientFamily = Field(description="The main ingredient")
+
+    with pytest.raises(ValueError, match="inconsistent key types"):
+        @tool_family(GardenParams)
+        class BadRecipeNested(WithImplementation):
+            """Write a recipe for the {dish}"""
+            extras: list[IngredientFamily] = Field(description="Additional ingredients")
+
+    with pytest.raises(ValueError, match="inconsistent key types"):
+        @tool_family(GardenParams)
+        class BadRecipeUndescribed(WithImplementation):
+            """Write a recipe for the {dish}"""
+            staple: IngredientFamily
+
+
 def test_injected_state_subscription_copies_doc():
     class Slice(TypedDict):
         n: int