diff --git a/graphcore/tools/schemas.py b/graphcore/tools/schemas.py index c42c713..236b76f 100644 --- a/graphcore/tools/schemas.py +++ b/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,59 +141,105 @@ 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) @@ -197,7 +247,15 @@ def wrapper(t: type[T]): 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 diff --git a/tests/test_tool_families.py b/tests/test_tool_families.py index 48aef89..58f858f 100644 --- a/tests/test_tool_families.py +++ b/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