Schema black magic for tool families

[jtoman]

Adds a new level of indirection for tool schemas: "families". These are tools that are generic in their behavior, but for LLM prompting purposes we want to vary the nouns in the schemas sent along. For example, a generic "edit the spec" tool might be present itself as "Edit your CVL spec file" to one agent, vs "Edit your foundry fuzz test" for another.

This is accomplished by using format strings in the Field descriptions and the doc string of any subclass with WithAsyncDependeicies or With*Implementation and then annotating the class with tool_family(Schema). Schema is the typename of a subclass of ToolFamilyParams. These classes are phantom types that provide strong typing for the template parameters.

The intended behavior is to declare:

class MyParams(ToolFamilyParams):
   field1: int
   field2: str
   # ...

and then annotate your tool with:

@tool_family(MyParams)

This transforms the tool into an object with a with_template(kwarg_sig) method which then returns (a subclass of) your original type. kwarg_sig is the kwarg signature derived from MyParams, equivalent to Unpack[MyParams] if MyParams was a TypedDict. The tool_family annotation will check that all format string params in your tool are defined by MyParams; the annotation will throw if this is not the case.

NB: The formatting is not transitive. If you do:

class SomethingElse(BaseModel):
    """
         My cool {thing}
    """
    # ...

@tool_family(...)
class MyTool(...):
    my_state: SomethingElse

you'll end up sending My cool {thing} to the llm. If actual transitive formatting is needed, we can add that later.

Here is a full example:

class MyToolParams(ToolFamilyParams):
    spec_noun: str
    component_noun: str

@tool_family(MyToolParams)
class MyTool(WithAsyncDependencies[str, str]):
    """
    A tool to check the validity of the {spec_noun} in the context of {component_noun}
    """

    l: str = Field(description="The {spec_noun} to check")

A typechecker rejects MyTool.bind(...) or MyTool.as_tool(...); the only MyTool supports is with_templates. You can instantiate this tool family with: MyTool.with_template(spec_noun="CVL File", component_noun="the smart contract").bind(...).as_tool(...)

[ericeil]

So I asked my trusty LLM what the , None means in this tuple, and it was like "oh that makes all fields optional." It wrote me this test to show that it fails:

import pytest
from pydantic import Field, ValidationError

from graphcore.tools.schemas import ToolFamilyParams, WithImplementation, tool_family


class CheckSpecParams(ToolFamilyParams):
    spec_noun: str
    component_noun: str


@tool_family(CheckSpecParams)
class CheckSpec(WithImplementation[str]):
    """A tool to check the validity of the {spec_noun} in the context of {component_noun}"""

    spec: str = Field(description="The {spec_noun} to check")

    def run(self) -> str:
        return self.spec


def test_with_template_keeps_required_fields_required():
    """with_template must not turn a required argument into an optional one.

    The current create_model call passes None as each overridden field's
    default (the Pydantic "make every field optional" pattern), so constructing
    the specialized tool with no arguments succeeds and spec is None.
    """
    specialized = CheckSpec.with_template(
        spec_noun="CVL File",
        component_noun="the smart contract",
    )

    assert specialized.model_fields["spec"].is_required()
    with pytest.raises(ValidationError):
        specialized()
    assert specialized(spec="rule.cvl").spec == "rule.cvl"

Suggested fix:

Suggested change

	for (k, v) in self._staged.model_fields.items():
	if not v.description:
	continue
	descr = v.asdict()
	new_attrs = {
	**descr["attributes"],
	"description": v.description.format(*args, **kwargs)
	}
	new_fields[k] = (Annotated[(v.annotation, *descr["metadata"], Field(**new_attrs))], None)
	for (k, v) in self._staged.model_fields.items():
	if not v.description or not _placeholders(v.description):
	continue
	descr = v.asdict()
	new_attrs = {
	**descr["attributes"],
	"description": v.description.format(*args, **kwargs)
	}
	if not v.is_required():
	if v.default_factory is not None:
	new_attrs.setdefault("default_factory", v.default_factory)
	else:
	new_attrs.setdefault("default", v.default)
	new_fields[k] = Annotated[(v.annotation, *descr["metadata"], Field(**new_attrs))]

[ericeil]

I mean, I guess you're not really intending this to be used this way. :)

[ericeil]

Looks good, just see the comment I left in case it actually matters.

[ericeil]

Test and fix look good; re-approving