diff --git a/graphcore/tools/schemas.py b/graphcore/tools/schemas.py index 3009049..79963c3 100644 --- a/graphcore/tools/schemas.py +++ b/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 \ No newline at end of file diff --git a/tests/test_tool_families.py b/tests/test_tool_families.py index 58f858f..7a170d5 100644 --- a/tests/test_tool_families.py +++ b/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