From 4191291145d3f8073ffc247f94f2f993db61f3bb Mon Sep 17 00:00:00 2001 From: Eric Eilebrecht Date: Thu, 20 Aug 2026 13:24:57 -0700 Subject: [PATCH 1/5] tool families: a rendered family parameter keeps its identity `family_param` promises a decorator that leaves the name usable as an annotation -- it returns `type[T]`, and the class it binds is a subclass of the decorated one. Two things then broke that promise at runtime while no checker could see it. `with_template` rendered onto `_wrapped`, so the class it produced was a sibling of the class the decorator bound, not a subtype of it. Annotate anything with that name -- graph state above all -- and every value a templated tool builds fails validation against it. Statically the two are one type, so nothing flags it; in langgraph the failure then lands under `loc=('state', ...)`, which is stripped as injected, leaving the model an empty error string it retries against forever. Neither the bound class nor a rendering of it admitted where it lived, either: `create_model` takes `__module__` from the calling frame, so both claimed this module. Restoring a value by importing its class -- what a checkpoint serializer does -- looked for it here, did not find it, and handed back a bare dict. A rendering's identity is the whole of what separates the two kinds of family, so they are now two types rather than one with a branch. A `_ToolFamily` handle is never a value's type: a rendering of it is the wrapped schema and stays where it is built, which no name resolves to, since the bound name holds the fieldless handle. A `_FamilyParam` renders onto itself and claims its own module, so a rendered value both validates as the bound class and comes back from a checkpoint as one. Co-Authored-By: Claude Opus 5 (1M context) --- graphcore/tools/schemas.py | 80 ++++++++++++++++++++++++++++++---- tests/test_tool_families.py | 87 ++++++++++++++++++++++++++++++++++++- 2 files changed, 157 insertions(+), 10 deletions(-) diff --git a/graphcore/tools/schemas.py b/graphcore/tools/schemas.py index 3009049..2815558 100644 --- a/graphcore/tools/schemas.py +++ b/graphcore/tools/schemas.py @@ -188,9 +188,29 @@ 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. What the two + variants below decide is that rendered class's identity: what it derives from, and where it + claims to live.""" + _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 _rendered_module(cls) -> str: + """The module a rendered schema claims. + + A serialized value names its class by module and class name, and is restored by importing + the one and looking the other up in it, so this decides what a rendered value comes back + as -- or whether it comes back as a value at all.""" + raise NotImplementedError + @classmethod def with_template(cls, *args: P.args, **kwargs: P.kwargs) -> T: assert cls._wrapped.__doc__ is not None @@ -218,30 +238,72 @@ def type_mapper( return create_model( cls._wrapped.__name__, __doc__=new_doc, - __base__=cast(T, cls._wrapped), + __base__=cast(T, cls._render_onto()), + __module__=cls._rendered_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 + + @override + @classmethod + def _rendered_module(cls) -> str: + # Stays where it is built, which no name resolves to: `t`'s own name in `t`'s module is + # this handle, and a rendering that claimed to be that would restore with no fields at all. + return __name__ + @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*. + + So a value a templated tool builds is an instance of the name the decorator bound, and + restoring one recovers that name -- the rendering itself is not importable, being built at + runtime, and this is the class it renders onto.""" + + @override + @classmethod + def _render_onto(cls) -> type[BaseModel]: + return cls + + @override + @classmethod + def _rendered_module(cls) -> str: + return cls.__module__ + @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 +354,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 +365,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..0ebdf9b 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,87 @@ 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" + + +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 + + +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 + + +def test_a_rendered_value_survives_a_checkpoint_round_trip(): + # The serializer restores a model by importing its class. A rendering is built at runtime and + # so is importable under no name; what it has to come back as is the class it renders onto. + serve = tool_family(RecipeParams)(ServeDish).with_template(dish="risotto") + plated = serve.model_validate({"portion": {"grams": 200}}) + + serde = JsonPlusSerializer() + (restored,) = serde.loads_typed(serde.dumps_typed([plated.portion])) + assert isinstance(restored, Portion), f"restored as {type(restored)}, not the bound class" + assert restored.grams == 200 + + +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" + + def test_inconsistent_key_types_rejected(): class GardenParams(ToolFamilyParams): dish: str From 2f07d69e90ba21e2d930c08827395c8aeae49f5d Mon Sep 17 00:00:00 2001 From: Eric Eilebrecht Date: Thu, 20 Aug 2026 14:14:18 -0700 Subject: [PATCH 2/5] Rebind family-param values at the tool/state boundary Stamping a rendering's __module__ as the bound class made JsonPlus look the bound name up, but the live object was still a different class. LangChain validates against the rendered schema and then calls the tool with those instances, so a second construction rejects a bound value typed as the rendering. as_tool, tool_state_update, and tool_output now rebuild rendered family-param instances as the decorator-bound class. What hits a checkpoint is importable; the rendering stays the LLM schema. --- graphcore/graph.py | 5 ++- graphcore/tools/schemas.py | 81 ++++++++++++++++++++----------------- tests/test_tool_families.py | 33 ++++++++++++--- 3 files changed, 74 insertions(+), 45 deletions(-) diff --git a/graphcore/graph.py b/graphcore/graph.py index 6cddf71..9c14579 100644 --- a/graphcore/graph.py +++ b/graphcore/graph.py @@ -38,6 +38,7 @@ from pydantic import BaseModel, ValidationError from .utils import ainvoke, invoke, current_prompt_tokens, get_token_usage from .summary import SummaryConfig, Summarization +from .tools.schemas import rebind_family_param_values logger = logging.getLogger(__name__) @@ -81,7 +82,7 @@ def tool_output(tool_call_id: str, res: dict) -> Command: Command that updates state with final results and a success message """ return Command(update={ - **res, + **rebind_family_param_values(res), "messages": [ToolMessage( tool_call_id=tool_call_id, content="Success" @@ -123,7 +124,7 @@ def tool_state_update( "messages": [ ToolMessage(tool_call_id=tool_call_id, content=content) ], - **state_diff + **rebind_family_param_values(state_diff) } return Command(update=update) diff --git a/graphcore/tools/schemas.py b/graphcore/tools/schemas.py index 2815558..b0f7aa0 100644 --- a/graphcore/tools/schemas.py +++ b/graphcore/tools/schemas.py @@ -38,6 +38,39 @@ def __class_getitem__(cls, params: Any) -> Any: class WithInjectedId(BaseModel): tool_call_id: Annotated[str, InjectedToolCallId] +def rebind_family_param_values(value: Any) -> Any: + """Replace rendered family-param instances with the class the decorator bound. + + A rendering is a runtime subclass used as the LLM schema. LangGraph restores a + model by importing its class, which only the bound name admits.""" + if isinstance(value, list): + return [rebind_family_param_values(v) for v in value] + if isinstance(value, tuple): + return tuple(rebind_family_param_values(v) for v in value) + if isinstance(value, dict): + return {k: rebind_family_param_values(v) for k, v in value.items()} + bound = getattr(type(value), "_bound", None) + if ( + isinstance(value, BaseModel) + and isinstance(bound, type) + and type(value) is not bound + and issubclass(type(value), bound) + ): + return bound.model_validate(value.model_dump()) + return value + + +def _tool_instance(cls: type[BaseModel], kwargs: dict[str, Any]) -> Any: + instance = cls(**kwargs) + for name, field in type(instance).model_fields.items(): + if InjectedState in field.metadata: + continue + object.__setattr__( + instance, name, rebind_family_param_values(getattr(instance, name)) + ) + return instance + + class WithImplementation(BaseModel, Generic[T_RES]): def run(self) -> T_RES: """Override this method to implement the tool logic.""" @@ -50,10 +83,8 @@ def as_tool( ) -> BaseTool: impl_method = getattr(cls, "run") - # Simple wrapper - just accept kwargs, instantiate model, call method def wrapper(**kwargs: Any) -> Any: - instance = cls(**kwargs) - return impl_method(instance) + return impl_method(_tool_instance(cls, kwargs)) return StructuredTool.from_function( func=wrapper, @@ -74,11 +105,8 @@ def as_tool( ) -> BaseTool: impl_method = getattr(cls, "run") - # Simple wrapper - just accept kwargs, instantiate model, call method async def wrapper(**kwargs: Any) -> Any: - instance = cls(**kwargs) - d = await impl_method(instance) - return d + return await impl_method(_tool_instance(cls, kwargs)) return StructuredTool.from_function( coroutine=wrapper, @@ -99,9 +127,8 @@ def __init__(self, ty: type[DEPS_BOUND], deps: object): def as_tool(self, name: str) -> BaseTool: impl_method = self._ty.run - # Simple wrapper - just accept kwargs, instantiate model, call method async def wrapper(**kwargs: Any) -> Any: - instance = self._ty(**kwargs) + instance = _tool_instance(self._ty, kwargs) tok = self._ty._dep_ctx.set(self.deps) try: d = await impl_method(instance) @@ -190,9 +217,8 @@ def map_type[T, U](t: Any, to_rewrite: type[T], f: Callable[[type[T]], type[U]]) 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. What the two - variants below decide is that rendered class's identity: what it derives from, and where it - claims to live.""" + :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]] @@ -202,15 +228,6 @@ def _render_onto(cls) -> type[BaseModel]: """The base a rendered schema derives from.""" raise NotImplementedError - @classmethod - def _rendered_module(cls) -> str: - """The module a rendered schema claims. - - A serialized value names its class by module and class name, and is restored by importing - the one and looking the other up in it, so this decides what a rendered value comes back - as -- or whether it comes back as a value at all.""" - raise NotImplementedError - @classmethod def with_template(cls, *args: P.args, **kwargs: P.kwargs) -> T: assert cls._wrapped.__doc__ is not None @@ -239,7 +256,6 @@ def type_mapper( cls._wrapped.__name__, __doc__=new_doc, __base__=cast(T, cls._render_onto()), - __module__=cls._rendered_module(), **new_fields ) @@ -254,13 +270,6 @@ class _ToolFamily[T: type[BaseModel], M: ToolFamilyParams, **P](_TemplatedTool[T def _render_onto(cls) -> type[BaseModel]: return cls._wrapped - @override - @classmethod - def _rendered_module(cls) -> str: - # Stays where it is built, which no name resolves to: `t`'s own name in `t`'s module is - # this handle, and a rendering that claimed to be that would restore with no fields at all. - return __name__ - @staticmethod def of[X: BaseModel, K: ToolFamilyParams, **R](t: type[X], m: type[K]) -> type["_ToolFamily[type[X], K, R]"]: clone = create_model( @@ -278,20 +287,17 @@ class _FamilyParam[T: type[BaseModel], M: ToolFamilyParams, **P](_TemplatedTool[ """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*. - So a value a templated tool builds is an instance of the name the decorator bound, and - restoring one recovers that name -- the rendering itself is not importable, being built at - runtime, and this is the class it renders onto.""" + Values written through :meth:`WithImplementation.as_tool` / :func:`rebind_family_param_values` + are this class, not the rendering. LangGraph restores a model by importing its class, which + only this (module-level) name admits.""" + + _bound: ClassVar[type[BaseModel]] @override @classmethod def _render_onto(cls) -> type[BaseModel]: return cls - @override - @classmethod - def _rendered_module(cls) -> str: - return cls.__module__ - @staticmethod def of[X: BaseModel, K: ToolFamilyParams, **R](t: type[X], m: type[K]) -> type[X]: clone = create_model( @@ -306,6 +312,7 @@ def of[X: BaseModel, K: ToolFamilyParams, **R](t: type[X], m: type[K]) -> type[ clone_narrowed = cast(type[_FamilyParam[type[X], K, R]], clone) clone_narrowed._wrapped = t clone_narrowed._key_type = m + clone_narrowed._bound = clone_narrowed return cast(type[X], clone_narrowed) diff --git a/tests/test_tool_families.py b/tests/test_tool_families.py index 0ebdf9b..0545a5b 100644 --- a/tests/test_tool_families.py +++ b/tests/test_tool_families.py @@ -12,8 +12,10 @@ from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer +from graphcore.graph import tool_state_update from graphcore.tools.schemas import ( WithImplementation, WithInjectedState, ToolFamilyParams, family_param, tool_family, + rebind_family_param_values, ) class TemplateArgValues(TypedDict): @@ -335,6 +337,7 @@ def test_family_param_value_validates_against_the_bound_name(): 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(): @@ -344,16 +347,33 @@ def test_family_param_lives_where_the_decorator_bound_it(): def test_a_rendered_value_survives_a_checkpoint_round_trip(): - # The serializer restores a model by importing its class. A rendering is built at runtime and - # so is importable under no name; what it has to come back as is the class it renders onto. - serve = tool_family(RecipeParams)(ServeDish).with_template(dish="risotto") - plated = serve.model_validate({"portion": {"grams": 200}}) + # JsonPlus names a model by module and class; a rendering is importable under no name. + # as_tool / tool_state_update rebind values to the bound class before they hit state. + 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 type(plated) is Portion serde = JsonPlusSerializer() - (restored,) = serde.loads_typed(serde.dumps_typed([plated.portion])) - assert isinstance(restored, Portion), f"restored as {type(restored)}, not the bound class" + (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 + cmd = tool_state_update("t1", "ok", portions=[raw]) + assert type(cmd.update["portions"][0]) is Portion + assert type(rebind_family_param_values(raw)) is Portion + def test_family_param_renderings_stay_independent(): stew = tool_family(RecipeParams)(ServeDish).with_template(dish="stew") @@ -374,6 +394,7 @@ def test_family_param_is_directly_renderable(): assert issubclass(portion, Portion) assert portion.__doc__ == "A portion of the curry" + assert type(rebind_family_param_values(portion(grams=3))) is Portion def test_inconsistent_key_types_rejected(): From 49cdb4baa9b46e5701090c577e6937c3de5e661d Mon Sep 17 00:00:00 2001 From: Eric Eilebrecht Date: Thu, 20 Aug 2026 14:16:39 -0700 Subject: [PATCH 3/5] Narrow Command.update before subscripting in the family-param test Pyright types Command.update as Any | None; assert it is present before indexing the rebound portions. --- tests/test_tool_families.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_tool_families.py b/tests/test_tool_families.py index 0545a5b..e4ec59f 100644 --- a/tests/test_tool_families.py +++ b/tests/test_tool_families.py @@ -371,6 +371,7 @@ def run(self) -> Portion: raw = rendered.model_validate({"portion": {"grams": 50}}).portion assert type(raw) is not Portion cmd = tool_state_update("t1", "ok", portions=[raw]) + assert cmd.update is not None assert type(cmd.update["portions"][0]) is Portion assert type(rebind_family_param_values(raw)) is Portion From 01474dc0809865eacb5a339ff7a04ef07359cbf1 Mon Sep 17 00:00:00 2001 From: Eric Eilebrecht Date: Thu, 20 Aug 2026 16:04:32 -0700 Subject: [PATCH 4/5] Stamp family-param renderings with the bound class's module create_model was taking __module__ from this file, so JsonPlus dumped a rendering as graphcore.tools.schemas.Portion, failed the lookup, and restored a dict. Rebind at as_tool / Command.update papered over that. A rendering now claims _render_onto()'s module, which for a family param is the class the decorator bound. Restore constructs that class. The live object can stay a subclass used as the LLM schema. --- graphcore/graph.py | 5 ++-- graphcore/tools/schemas.py | 56 +++++++++---------------------------- tests/test_tool_families.py | 27 +++++++++++------- 3 files changed, 32 insertions(+), 56 deletions(-) diff --git a/graphcore/graph.py b/graphcore/graph.py index 9c14579..6cddf71 100644 --- a/graphcore/graph.py +++ b/graphcore/graph.py @@ -38,7 +38,6 @@ from pydantic import BaseModel, ValidationError from .utils import ainvoke, invoke, current_prompt_tokens, get_token_usage from .summary import SummaryConfig, Summarization -from .tools.schemas import rebind_family_param_values logger = logging.getLogger(__name__) @@ -82,7 +81,7 @@ def tool_output(tool_call_id: str, res: dict) -> Command: Command that updates state with final results and a success message """ return Command(update={ - **rebind_family_param_values(res), + **res, "messages": [ToolMessage( tool_call_id=tool_call_id, content="Success" @@ -124,7 +123,7 @@ def tool_state_update( "messages": [ ToolMessage(tool_call_id=tool_call_id, content=content) ], - **rebind_family_param_values(state_diff) + **state_diff } return Command(update=update) diff --git a/graphcore/tools/schemas.py b/graphcore/tools/schemas.py index b0f7aa0..c6094b2 100644 --- a/graphcore/tools/schemas.py +++ b/graphcore/tools/schemas.py @@ -38,39 +38,6 @@ def __class_getitem__(cls, params: Any) -> Any: class WithInjectedId(BaseModel): tool_call_id: Annotated[str, InjectedToolCallId] -def rebind_family_param_values(value: Any) -> Any: - """Replace rendered family-param instances with the class the decorator bound. - - A rendering is a runtime subclass used as the LLM schema. LangGraph restores a - model by importing its class, which only the bound name admits.""" - if isinstance(value, list): - return [rebind_family_param_values(v) for v in value] - if isinstance(value, tuple): - return tuple(rebind_family_param_values(v) for v in value) - if isinstance(value, dict): - return {k: rebind_family_param_values(v) for k, v in value.items()} - bound = getattr(type(value), "_bound", None) - if ( - isinstance(value, BaseModel) - and isinstance(bound, type) - and type(value) is not bound - and issubclass(type(value), bound) - ): - return bound.model_validate(value.model_dump()) - return value - - -def _tool_instance(cls: type[BaseModel], kwargs: dict[str, Any]) -> Any: - instance = cls(**kwargs) - for name, field in type(instance).model_fields.items(): - if InjectedState in field.metadata: - continue - object.__setattr__( - instance, name, rebind_family_param_values(getattr(instance, name)) - ) - return instance - - class WithImplementation(BaseModel, Generic[T_RES]): def run(self) -> T_RES: """Override this method to implement the tool logic.""" @@ -84,7 +51,8 @@ def as_tool( impl_method = getattr(cls, "run") def wrapper(**kwargs: Any) -> Any: - return impl_method(_tool_instance(cls, kwargs)) + instance = cls(**kwargs) + return impl_method(instance) return StructuredTool.from_function( func=wrapper, @@ -106,7 +74,9 @@ def as_tool( impl_method = getattr(cls, "run") async def wrapper(**kwargs: Any) -> Any: - return await impl_method(_tool_instance(cls, kwargs)) + instance = cls(**kwargs) + d = await impl_method(instance) + return d return StructuredTool.from_function( coroutine=wrapper, @@ -128,7 +98,7 @@ def as_tool(self, name: str) -> BaseTool: impl_method = self._ty.run async def wrapper(**kwargs: Any) -> Any: - instance = _tool_instance(self._ty, kwargs) + instance = self._ty(**kwargs) tok = self._ty._dep_ctx.set(self.deps) try: d = await impl_method(instance) @@ -252,10 +222,14 @@ 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._render_onto()), + __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 ) @@ -287,11 +261,8 @@ class _FamilyParam[T: type[BaseModel], M: ToolFamilyParams, **P](_TemplatedTool[ """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*. - Values written through :meth:`WithImplementation.as_tool` / :func:`rebind_family_param_values` - are this class, not the rendering. LangGraph restores a model by importing its class, which - only this (module-level) name admits.""" - - _bound: ClassVar[type[BaseModel]] + 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 @@ -312,7 +283,6 @@ def of[X: BaseModel, K: ToolFamilyParams, **R](t: type[X], m: type[K]) -> type[ clone_narrowed = cast(type[_FamilyParam[type[X], K, R]], clone) clone_narrowed._wrapped = t clone_narrowed._key_type = m - clone_narrowed._bound = clone_narrowed return cast(type[X], clone_narrowed) diff --git a/tests/test_tool_families.py b/tests/test_tool_families.py index e4ec59f..7a170d5 100644 --- a/tests/test_tool_families.py +++ b/tests/test_tool_families.py @@ -12,10 +12,8 @@ from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer -from graphcore.graph import tool_state_update from graphcore.tools.schemas import ( WithImplementation, WithInjectedState, ToolFamilyParams, family_param, tool_family, - rebind_family_param_values, ) class TemplateArgValues(TypedDict): @@ -327,6 +325,7 @@ def test_family_param_renders_a_subtype_of_the_bound_name(): 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(): @@ -345,10 +344,16 @@ def test_family_param_lives_where_the_decorator_bound_it(): 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 is importable under no name. - # as_tool / tool_state_update rebind values to the bound class before they hit state. + # 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") @@ -357,7 +362,8 @@ def run(self) -> Portion: tool = tool_family(RecipeParams)(Plate).with_template(dish="risotto").as_tool("plate") plated = tool.invoke({"portion": {"grams": 200}}) - assert type(plated) is Portion + assert isinstance(plated, Portion) + assert type(plated) is not Portion serde = JsonPlusSerializer() (restored,) = serde.loads_typed(serde.dumps_typed([plated])) @@ -370,10 +376,9 @@ def run(self) -> Portion: rendered = tool_family(RecipeParams)(ServeDish).with_template(dish="stew") raw = rendered.model_validate({"portion": {"grams": 50}}).portion assert type(raw) is not Portion - cmd = tool_state_update("t1", "ok", portions=[raw]) - assert cmd.update is not None - assert type(cmd.update["portions"][0]) is Portion - assert type(rebind_family_param_values(raw)) is 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(): @@ -395,7 +400,9 @@ def test_family_param_is_directly_renderable(): assert issubclass(portion, Portion) assert portion.__doc__ == "A portion of the curry" - assert type(rebind_family_param_values(portion(grams=3))) is Portion + serde = JsonPlusSerializer() + restored = serde.loads_typed(serde.dumps_typed(portion(grams=3))) + assert type(restored) is Portion def test_inconsistent_key_types_rejected(): From f10cdd3fe03f6d244d14a1109831143668e05e80 Mon Sep 17 00:00:00 2001 From: Eric Eilebrecht Date: Thu, 20 Aug 2026 16:18:50 -0700 Subject: [PATCH 5/5] Undo unneeded comment removal --- graphcore/tools/schemas.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/graphcore/tools/schemas.py b/graphcore/tools/schemas.py index c6094b2..79963c3 100644 --- a/graphcore/tools/schemas.py +++ b/graphcore/tools/schemas.py @@ -50,6 +50,7 @@ def as_tool( ) -> BaseTool: impl_method = getattr(cls, "run") + # Simple wrapper - just accept kwargs, instantiate model, call method def wrapper(**kwargs: Any) -> Any: instance = cls(**kwargs) return impl_method(instance) @@ -73,6 +74,7 @@ def as_tool( ) -> BaseTool: impl_method = getattr(cls, "run") + # Simple wrapper - just accept kwargs, instantiate model, call method async def wrapper(**kwargs: Any) -> Any: instance = cls(**kwargs) d = await impl_method(instance) @@ -97,6 +99,7 @@ def __init__(self, ty: type[DEPS_BOUND], deps: object): def as_tool(self, name: str) -> BaseTool: impl_method = self._ty.run + # Simple wrapper - just accept kwargs, instantiate model, call method async def wrapper(**kwargs: Any) -> Any: instance = self._ty(**kwargs) tok = self._ty._dep_ctx.set(self.deps)