Skip to content

feat: add tool decorator #387

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.

Sign up for GitHub

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

Open
akihikokuroda wants to merge 6 commits into
base: main
Choose a base branch
from
+539 −6
Open

feat: add tool decorator #387

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
e846668
add tool decorator
akihikokuroda Jan 30, 2026
9c0fda4
review comments
akihikokuroda Feb 2, 2026
4d384e6
fix merge error
akihikokuroda Feb 3, 2026
e564cdb
fix lint error
akihikokuroda Feb 3, 2026
86440b4
review comments
akihikokuroda Feb 4, 2026
4095165
review comment
akihikokuroda Feb 4, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion docs/examples/tools/interpreter_example.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
# pytest: ollama, llm

from mellea import MelleaSession, start_session
from mellea.backends import ModelOption
from mellea.backends import ModelOption, tool
from mellea.stdlib.requirements import tool_arg_validator, uses_tool
from mellea.stdlib.tools import code_interpreter, local_code_interpreter

Expand Down
137 changes: 137 additions & 0 deletions docs/examples/tools/tool_decorator_example.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,137 @@
# pytest: ollama, llm
"""Example demonstrating the @tool decorator for cleaner tool definitions."""

import ast

from mellea import start_session
from mellea.backends import ModelOption, tool


# Define tools using the @tool decorator - much cleaner than MelleaTool.from_callable()
@tool
def get_weather(location: str, days: int = 1) -> dict:
"""Get weather forecast for a location.

Args:
location: City name
days: Number of days to forecast (default: 1)
"""
# Mock implementation
return {"location": location, "days": days, "forecast": "sunny", "temperature": 72}


@tool
def search_web(query: str, max_results: int = 5) -> list[str]:
"""Search the web for information.

Args:
query: Search query
max_results: Maximum number of results to return
"""
# Mock implementation
return [f"Result {i + 1} for '{query}'" for i in range(max_results)]


@tool(name="calculator")
def calculate(expression: str) -> str:
"""Evaluate a mathematical expression.

Args:
expression: Mathematical expression to evaluate
"""
try:
# Use ast.literal_eval for safe evaluation of simple expressions
result = ast.literal_eval(expression)
return f"Result: {result}"
except Exception as e:
return f"Error: {e!s}"


def example_basic_usage():
"""Example 1: Basic usage with decorated tools."""
print("\n=== Example 1: Basic Tool Usage ===")

# Without the decorator, you can add tools using:
# tools = [MelleaTool.from_callable(get_weather), MelleaTool.from_callable(search_web)]

# Now you can just pass the decorated functions directly to model_options
# Example: model_options={ModelOption.TOOLS: [get_weather, search_web, calculate]}

# The decorated tools must be called using .run()
weather = get_weather.run("Boston", days=3)
print(f"Tool call via .run(): {weather}")

# And they have tool properties
print(f"Tool name: {get_weather.name}")
print(f"Tool has JSON schema: {'function' in get_weather.as_json_tool}")


def example_with_llm():
"""Example 2: Using decorated tools with an LLM."""
print("\n=== Example 2: Using Tools with LLM ===")

m = start_session()

# Pass decorated tools directly - no wrapping needed!
response = m.instruct(
description="What's the weather like in San Francisco?",
model_options={ModelOption.TOOLS: [get_weather, search_web]},
)

print(f"Response: {response}")


def example_custom_name():
"""Example 3: Using custom tool names."""
print("\n=== Example 3: Custom Tool Names ===")

# The calculator tool was decorated with @tool(name="calculator")
# So its name is "calculator" instead of "calculate"
print("Function name: calculate")
print(f"Tool name: {calculate.name}")

# Must use .run() to invoke
result = calculate.run("2 + 2")
print(f"Result: {result}")


def example_comparison():
"""Example 4: Comparison of old vs new approach."""
print("\n=== Example 4: Old vs New Approach ===")

# OLD APPROACH (still works, but verbose):
from mellea.backends.tools import MelleaTool

def old_style_tool(x: int) -> int:
"""Old style tool.

Args:
x: Input value
"""
return x * 2

old_tool = MelleaTool.from_callable(old_style_tool)
print(f"Old approach - tool name: {old_tool.name}")

# NEW APPROACH (cleaner):
@tool
def new_style_tool(x: int) -> int:
"""New style tool.

Args:
x: Input value
"""
return x * 2

print(f"New approach - tool name: {new_style_tool.name}")

# Both can be used together in a tools list
tools = [old_tool, new_style_tool, get_weather]
print(f"Mixed tools list: {[t.name for t in tools]}")


if __name__ == "__main__":
example_basic_usage()
# example_with_llm() # Uncomment to test with actual LLM
example_custom_name()
example_comparison()
8 changes: 7 additions & 1 deletion docs/tutorial.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1347,9 +1347,15 @@ For examples on adding tools to the template representation of a component, see
Here's an example of adding a tool through model options. This can be useful when you want to add a tool like web search that should almost always be available:
```python
import mellea
from mellea.backends import ModelOption
from mellea.backends import ModelOption, tool

@tool
def web_search(query: str) -> str:
"""Search the web for information.

Args:
query: The search query
"""
...

m = mellea.start_session()
Expand Down
3 changes: 3 additions & 0 deletions mellea/backends/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,12 +6,15 @@
from .cache import SimpleLRUCache
from .model_ids import ModelIdentifier
from .model_options import ModelOption
from .tools import MelleaTool, tool

__all__ = [
"Backend",
"BaseModelSubclass",
"FormatterBackend",
"MelleaTool",
"ModelIdentifier",
"ModelOption",
"SimpleLRUCache",
"tool",
]
86 changes: 82 additions & 4 deletions mellea/backends/tools.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
import re
from collections import defaultdict
from collections.abc import Callable, Generator, Iterable, Mapping, Sequence
from typing import Any, Literal
from typing import Any, Literal, overload

from pydantic import BaseModel, ConfigDict, Field

Expand Down Expand Up @@ -91,10 +91,84 @@ def from_callable(cls, func: Callable, name: str | None = None):
return MelleaTool(tool_name, tool_call, as_json)


@overload
def tool(func: Callable) -> MelleaTool: ...
Comment on lines +94 to +95
Copy link
Contributor

@jakelorocco jakelorocco Feb 4, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we add , *, name: str | None = None to this function definition so that you can still change the tool name when you call it on a function directly:

def new_tool(): ...

differently_named_tool = tool(new_tool, name="different_name")



@overload
def tool(*, name: str | None = None) -> Callable[[Callable], MelleaTool]: ...


def tool(
func: Callable | None = None, name: str | None = None
) -> MelleaTool | Callable[[Callable], MelleaTool]:
Comment on lines 102 to 104
Copy link
Contributor

@jakelorocco jakelorocco Feb 3, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This doesn't work because of your typing. If you do something like:

@tool
def tooling():
    ...

tooling()

You get type hinting errors on tooling() since it thinks it's a regular MelleaTool. For now, I think we should just return a regular MelleaTool and force it to be called with <tool>.run() like the other tools created from from_callable. If users get annoyed, we can add this functionality here.

Copy link
Member Author

@akihikokuroda akihikokuroda Feb 4, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I made changes. Thanks.

"""Decorator to mark a function as a Mellea tool.

This decorator wraps a function to make it usable as a tool without
requiring explicit MelleaTool.from_callable() calls. The decorated
function returns a MelleaTool instance that must be called via .run().

Args:
func: The function to decorate (when used without arguments)
name: Optional custom name for the tool (defaults to function name)

Returns:
A MelleaTool instance. Use .run() to invoke the tool.
The returned object passes isinstance(result, MelleaTool) checks.

Examples:
Basic usage:
>>> @tool
... def get_weather(location: str, days: int = 1) -> dict:
... '''Get weather forecast.
...
... Args:
... location: City name
... days: Number of days to forecast
... '''
... return {"location": location, "forecast": "sunny"}
>>>
>>> # The decorated function IS a MelleaTool
>>> isinstance(get_weather, MelleaTool) # True
>>>
>>> # Can be used directly in tools list (no extraction needed)
>>> tools = [get_weather]
>>>
>>> # Must use .run() to invoke the tool
>>> result = get_weather.run(location="Boston")

With custom name (as decorator):
>>> @tool(name="weather_api")
... def get_weather(location: str) -> dict:
... return {"location": location}
>>>
>>> result = get_weather.run(location="New York")

With custom name (as function):
>>> def new_tool(): ...
>>> differently_named_tool = tool(new_tool, name="different_name")
"""

def decorator(f: Callable) -> MelleaTool:
# Simply return the base MelleaTool instance
return MelleaTool.from_callable(f, name=name)

# Handle both @tool and @tool() syntax
if func is None:
# Called with arguments: @tool(name="custom")
return decorator
else:
# Called without arguments: @tool
return decorator(func)


def add_tools_from_model_options(
tools_dict: dict[str, AbstractMelleaTool], model_options: dict[str, Any]
):
"""If model_options has tools, add those tools to the tools_dict."""
"""If model_options has tools, add those tools to the tools_dict.

Accepts MelleaTool instances or @tool decorated functions.
"""
model_opts_tools = model_options.get(ModelOption.TOOLS, None)
if model_opts_tools is None:
return
Expand All @@ -110,16 +184,20 @@ def add_tools_from_model_options(
assert isinstance(tool_name, str), (
f"If ModelOption.TOOLS is a dict, it must be a dict of [str, Tool]; found {type(tool_name)} as the key instead"
)
assert isinstance(tool_instance, MelleaTool), (
assert isinstance(tool_instance, AbstractMelleaTool), (
f"If ModelOption.TOOLS is a dict, it must be a dict of [str, Tool]; found {type(tool_instance)} as the value instead"
)
tools_dict[tool_name] = tool_instance
else:
# Handle any other iterable / list here.
for tool_instance in model_opts_tools:
assert isinstance(tool_instance, MelleaTool), (
assert isinstance(tool_instance, AbstractMelleaTool), (
f"If ModelOption.TOOLS is a list, it must be a list of Tool; found {type(tool_instance)}"
)
# MelleaTool (and subclasses like CallableMelleaTool) have a name attribute
assert isinstance(tool_instance, MelleaTool), (
f"Tool must be a MelleaTool instance with a name attribute; found {type(tool_instance)}"
)
tools_dict[tool_instance.name] = tool_instance


Expand Down
Loading