Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Allow custom function_schema definition for Tool dataclass #878

Open
jonchun opened this issue Feb 9, 2025 · 1 comment · May be fixed by #881
Open

Allow custom function_schema definition for Tool dataclass #878

jonchun opened this issue Feb 9, 2025 · 1 comment · May be fixed by #881
Labels
Feature request New feature request

Comments

@jonchun
Copy link

jonchun commented Feb 9, 2025

Currently, the tool implementation assumes that all tools for agents are implemented at compile-time as pre-defined Python functions. It automatically generates a FunctionSchema based off the input function. If we supported a custom function_schema, it would allow for more advanced used cases such as defining tools at runtime (e.g. MCP)

I've attempted a naive fix for this as follows:

@dataclass(init=False)
class Tool(Generic[AgentDepsT]):
     ...
    def __init__(
        self,
        function: ToolFuncEither[AgentDepsT],
        *,
        takes_ctx: bool | None = None,
        max_retries: int | None = None,
        name: str | None = None,
        description: str | None = None,
        prepare: ToolPrepareFunc[AgentDepsT] | None = None,
        docstring_format: DocstringFormat = 'auto',
        require_parameter_descriptions: bool = False,
        function_schema: _pydantic.FunctionSchema | None = None,
    ):
        f = function_schema or _pydantic.function_schema(
            function, takes_ctx, docstring_format, require_parameter_descriptions
        )
        ...

Unfortunately this keeps getting wiped when registering the tool with the Agent. I've narrowed it down to the replace functionality here. I tried a few tweaks such as attempting to make function_schema a private/public attribute of Tool and wasn't able to get it to work correctly without breaking a bunch of tests. I haven't had a chance to troubleshoot further than this.

class Agent(Generic[AgentDepsT, ResultDataT]):
    ...
    def _register_tool(self, tool: Tool[AgentDepsT]) -> None:
        """Private utility to register a tool instance."""
        if tool.max_retries is None:
            # noinspection PyTypeChecker
            tool = dataclasses.replace(tool, max_retries=self._default_retries)
@jonchun
Copy link
Author

jonchun commented Feb 10, 2025

Resolved the above issue -- I was being silly and accidentally copy-pasted = field(init=False) so it wasn't automatically passing the parameter with dataclasses.replace(). Opened a tentative PR for this feature.

@sydney-runkle sydney-runkle added the Feature request New feature request label Feb 13, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
Feature request New feature request
Projects
None yet
Development

Successfully merging a pull request may close this issue.

2 participants