From fb465723aa4f411ad6c23d770542569a318763f7 Mon Sep 17 00:00:00 2001 From: "Paul S. Schweigert" Date: Thu, 23 Apr 2026 22:55:25 -0400 Subject: [PATCH 1/5] fix: add tool is_internal flag and populate context_view in pre-execute hooks Signed-off-by: Paul S. Schweigert --- docs/examples/plugins/tool_hooks.py | 5 +++ mellea/backends/tools.py | 16 +++++-- mellea/core/base.py | 4 ++ mellea/stdlib/components/react.py | 2 +- mellea/stdlib/functional.py | 1 + test/plugins/test_hook_call_sites.py | 23 ++++++++++ test/plugins/test_internal_tools.py | 43 +++++++++++++++++++ .../stdlib/frameworks/test_react_framework.py | 8 +++- 8 files changed, 96 insertions(+), 6 deletions(-) create mode 100644 test/plugins/test_internal_tools.py diff --git a/docs/examples/plugins/tool_hooks.py b/docs/examples/plugins/tool_hooks.py index fcd2e2fad..b1b4cd98e 100644 --- a/docs/examples/plugins/tool_hooks.py +++ b/docs/examples/plugins/tool_hooks.py @@ -150,6 +150,11 @@ def parse_factor(): @hook(HookType.TOOL_PRE_INVOKE, mode=PluginMode.CONCURRENT, priority=5) async def enforce_tool_allowlist(payload, _): """Block any tool not on the explicit allow list.""" + # Internal tools (like ReAct's final_answer) are always permitted. + if payload.model_tool_call.func.is_internal: + log.info("[allowlist] permitted internal tool=%r", payload.model_tool_call.name) + return None + tool_name = payload.model_tool_call.name if tool_name not in ALLOWED_TOOLS: log.warning( diff --git a/mellea/backends/tools.py b/mellea/backends/tools.py index 3dd056512..da1d57b16 100644 --- a/mellea/backends/tools.py +++ b/mellea/backends/tools.py @@ -34,6 +34,7 @@ class MelleaTool(AbstractMelleaTool): tool_call (Callable): The underlying Python callable to invoke when the tool is run. as_json_tool (dict[str, Any]): The OpenAI-compatible JSON schema dict describing the tool's parameters. + is_internal (bool): Whether the tool is internal to Mellea. Defaults to ``False``. """ @@ -41,14 +42,20 @@ class MelleaTool(AbstractMelleaTool): # Our ModelToolCall is the class that has a reference to the tool and actually calls with arguments name: str + is_internal: bool _as_json_tool: dict[str, Any] _call_tool: Callable[..., Any] def __init__( - self, name: str, tool_call: Callable, as_json_tool: dict[str, Any] + self, + name: str, + tool_call: Callable, + as_json_tool: dict[str, Any], + is_internal: bool = False, ) -> None: """Initialize the tool with a name, tool call and as_json_tool dict.""" self.name = name + self.is_internal = is_internal self._as_json_tool = as_json_tool self._call_tool = tool_call @@ -172,7 +179,9 @@ def tool_call(*args, **kwargs): ) from e @classmethod - def from_callable(cls, func: Callable, name: str | None = None) -> "MelleaTool": + def from_callable( + cls, func: Callable, name: str | None = None, is_internal: bool = False + ) -> "MelleaTool": """Create a MelleaTool from a plain Python callable. Introspects the callable's signature and docstring to build an @@ -181,6 +190,7 @@ def from_callable(cls, func: Callable, name: str | None = None) -> "MelleaTool": Args: func (Callable): The Python callable to wrap as a tool. name (str | None): Optional name override; defaults to ``func.__name__``. + is_internal (bool): Whether the tool is internal to Mellea. Defaults to ``False``. Returns: MelleaTool: A Mellea tool wrapping the callable. @@ -191,7 +201,7 @@ def from_callable(cls, func: Callable, name: str | None = None) -> "MelleaTool": exclude_none=True ) tool_call = func - return MelleaTool(tool_name, tool_call, as_json) + return MelleaTool(tool_name, tool_call, as_json, is_internal=is_internal) @overload diff --git a/mellea/core/base.py b/mellea/core/base.py index a3b40b97e..902861bf5 100644 --- a/mellea/core/base.py +++ b/mellea/core/base.py @@ -954,6 +954,7 @@ class AbstractMelleaTool(abc.ABC): Attributes: name (str): The unique name used to identify the tool in JSON descriptions and tool-call dispatch. + is_internal (bool): Whether the tool is internal to Mellea and should be bypassed by certain security policies. as_json_tool (dict[str, Any]): A JSON-serialisable description of the tool, compatible with the function-calling schemas expected by supported inference backends. """ @@ -961,6 +962,9 @@ class AbstractMelleaTool(abc.ABC): name: str """Name of the tool.""" + is_internal: bool = False + """Whether the tool is internal to Mellea.""" + @abc.abstractmethod def run(self, *args: Any, **kwargs: Any) -> Any: """Executes the tool with the provided arguments and returns the result. diff --git a/mellea/stdlib/components/react.py b/mellea/stdlib/components/react.py index 8f08fe8a0..6d0234cfc 100644 --- a/mellea/stdlib/components/react.py +++ b/mellea/stdlib/components/react.py @@ -71,7 +71,7 @@ def format_for_llm(self) -> TemplateRepresentation: ) finalizer_tool = MelleaTool.from_callable( - _mellea_finalize_tool, MELLEA_FINALIZER_TOOL + _mellea_finalize_tool, MELLEA_FINALIZER_TOOL, is_internal=True ) tools[MELLEA_FINALIZER_TOOL] = finalizer_tool diff --git a/mellea/stdlib/functional.py b/mellea/stdlib/functional.py index f4bdb2371..d23866a7e 100644 --- a/mellea/stdlib/functional.py +++ b/mellea/stdlib/functional.py @@ -589,6 +589,7 @@ async def aact( pre_exec_payload = ComponentPreExecutePayload( component_type=_component_type_name, action=action, + context_view=context.view_for_generation(), requirements=requirements or [], model_options=model_options or {}, format=format, diff --git a/test/plugins/test_hook_call_sites.py b/test/plugins/test_hook_call_sites.py index a8689380f..428b5e181 100644 --- a/test/plugins/test_hook_call_sites.py +++ b/test/plugins/test_hook_call_sites.py @@ -264,6 +264,29 @@ async def recorder(payload: Any, ctx: Any) -> Any: await aact(action, ctx, backend, strategy=None) assert observed[0].component_type == "Instruction" + async def test_component_pre_execute_has_context_view(self) -> None: + """COMPONENT_PRE_EXECUTE payload.context_view is populated.""" + from mellea.stdlib.components import Instruction + from mellea.stdlib.functional import aact + from mellea.stdlib.context import ChatContext + + observed: list[Any] = [] + + @hook("component_pre_execute") + async def recorder(payload: Any, ctx: Any) -> Any: + observed.append(payload) + return None + + register(recorder) + backend = _MockBackend() + ctx = ChatContext().add(CBlock("previous turn")) + action = Instruction("Context check") + + await aact(action, ctx, backend, strategy=None) + assert observed[0].context_view is not None + assert len(observed[0].context_view) == 1 + assert observed[0].context_view[0].value == "previous turn" + async def test_component_post_success_fires_in_aact(self) -> None: """COMPONENT_POST_SUCCESS fires in aact() after successful generation.""" from mellea.stdlib.components import Instruction diff --git a/test/plugins/test_internal_tools.py b/test/plugins/test_internal_tools.py new file mode 100644 index 000000000..ca9ea3409 --- /dev/null +++ b/test/plugins/test_internal_tools.py @@ -0,0 +1,43 @@ +import pytest +from mellea.backends.tools import MelleaTool +from mellea.core.base import ModelToolCall +from mellea.plugins import hook, HookType, PluginMode, block, register +from mellea.stdlib.functional import _acall_tools +from mellea.core.backend import Backend +from mellea.core.base import ModelOutputThunk, GenerateLog +from typing import Any + +class _MockBackend(Backend): + async def _generate_from_context(self, *args, **kwargs): + pass + async def generate_from_raw(self, *args, **kwargs): + pass + +@pytest.mark.asyncio +async def test_tool_is_internal_accessible_in_hook(): + """Verify that is_internal flag is accessible in TOOL_PRE_INVOKE hook.""" + + observed_is_internal = [] + + @hook(HookType.TOOL_PRE_INVOKE) + async def checker(payload, _): + observed_is_internal.append(payload.model_tool_call.func.is_internal) + return None + + register(checker) + + # Create an internal tool + internal_tool = MelleaTool.from_callable(lambda: "ok", name="internal", is_internal=True) + external_tool = MelleaTool.from_callable(lambda: "ok", name="external", is_internal=False) + + tc_internal = ModelToolCall(name="internal", func=internal_tool, args={}) + tc_external = ModelToolCall(name="external", func=external_tool, args={}) + + mot = ModelOutputThunk(value="", tool_calls={"internal": tc_internal, "external": tc_external}) + + await _acall_tools(mot, _MockBackend()) + + assert observed_is_internal == [True, False] + +if __name__ == "__main__": + pytest.main([__file__]) diff --git a/test/stdlib/frameworks/test_react_framework.py b/test/stdlib/frameworks/test_react_framework.py index e121a91f5..5136772ba 100644 --- a/test/stdlib/frameworks/test_react_framework.py +++ b/test/stdlib/frameworks/test_react_framework.py @@ -98,7 +98,9 @@ def _fn() -> str: def _final_answer_call(answer: str = "42") -> _ScriptedTurn: """Script a turn where the model calls final_answer with real arg flow.""" - tool = MelleaTool.from_callable(_mellea_finalize_tool, MELLEA_FINALIZER_TOOL) + tool = MelleaTool.from_callable( + _mellea_finalize_tool, MELLEA_FINALIZER_TOOL, is_internal=True + ) tc = ModelToolCall(name=MELLEA_FINALIZER_TOOL, func=tool, args={"answer": answer}) return _ScriptedTurn(value="", tool_calls={MELLEA_FINALIZER_TOOL: tc}) @@ -205,7 +207,9 @@ async def test_react_format_triggers_second_generation(): async def test_react_final_answer_with_extra_tool_rejected(): """final_answer alongside another tool in the same turn triggers assertion.""" search = _make_tool("search", "found it") - finalizer = MelleaTool.from_callable(_mellea_finalize_tool, MELLEA_FINALIZER_TOOL) + finalizer = MelleaTool.from_callable( + _mellea_finalize_tool, MELLEA_FINALIZER_TOOL, is_internal=True + ) both = { "search": ModelToolCall(name="search", func=search, args={}), MELLEA_FINALIZER_TOOL: ModelToolCall( From a707ddff5c281015c28da76fbb3ab1b98c9af81a Mon Sep 17 00:00:00 2001 From: "Paul S. Schweigert" Date: Fri, 24 Apr 2026 06:50:13 -0400 Subject: [PATCH 2/5] linting Signed-off-by: Paul S. Schweigert --- test/plugins/test_hook_call_sites.py | 2 +- test/plugins/test_internal_tools.py | 43 +++++++++++++++++----------- 2 files changed, 28 insertions(+), 17 deletions(-) diff --git a/test/plugins/test_hook_call_sites.py b/test/plugins/test_hook_call_sites.py index 428b5e181..cb893a900 100644 --- a/test/plugins/test_hook_call_sites.py +++ b/test/plugins/test_hook_call_sites.py @@ -267,8 +267,8 @@ async def recorder(payload: Any, ctx: Any) -> Any: async def test_component_pre_execute_has_context_view(self) -> None: """COMPONENT_PRE_EXECUTE payload.context_view is populated.""" from mellea.stdlib.components import Instruction - from mellea.stdlib.functional import aact from mellea.stdlib.context import ChatContext + from mellea.stdlib.functional import aact observed: list[Any] = [] diff --git a/test/plugins/test_internal_tools.py b/test/plugins/test_internal_tools.py index ca9ea3409..ec70bd79b 100644 --- a/test/plugins/test_internal_tools.py +++ b/test/plugins/test_internal_tools.py @@ -1,43 +1,54 @@ +from typing import Any + import pytest + from mellea.backends.tools import MelleaTool -from mellea.core.base import ModelToolCall -from mellea.plugins import hook, HookType, PluginMode, block, register -from mellea.stdlib.functional import _acall_tools from mellea.core.backend import Backend -from mellea.core.base import ModelOutputThunk, GenerateLog -from typing import Any +from mellea.core.base import GenerateLog, ModelOutputThunk, ModelToolCall +from mellea.plugins import HookType, PluginMode, block, hook, register +from mellea.stdlib.functional import _acall_tools + class _MockBackend(Backend): async def _generate_from_context(self, *args, **kwargs): pass + async def generate_from_raw(self, *args, **kwargs): pass + @pytest.mark.asyncio async def test_tool_is_internal_accessible_in_hook(): """Verify that is_internal flag is accessible in TOOL_PRE_INVOKE hook.""" - + observed_is_internal = [] - + @hook(HookType.TOOL_PRE_INVOKE) async def checker(payload, _): observed_is_internal.append(payload.model_tool_call.func.is_internal) return None - + register(checker) - + # Create an internal tool - internal_tool = MelleaTool.from_callable(lambda: "ok", name="internal", is_internal=True) - external_tool = MelleaTool.from_callable(lambda: "ok", name="external", is_internal=False) - + internal_tool = MelleaTool.from_callable( + lambda: "ok", name="internal", is_internal=True + ) + external_tool = MelleaTool.from_callable( + lambda: "ok", name="external", is_internal=False + ) + tc_internal = ModelToolCall(name="internal", func=internal_tool, args={}) tc_external = ModelToolCall(name="external", func=external_tool, args={}) - - mot = ModelOutputThunk(value="", tool_calls={"internal": tc_internal, "external": tc_external}) - + + mot = ModelOutputThunk( + value="", tool_calls={"internal": tc_internal, "external": tc_external} + ) + await _acall_tools(mot, _MockBackend()) - + assert observed_is_internal == [True, False] + if __name__ == "__main__": pytest.main([__file__]) From 5141340305f49016e1dbc81f0698c0f1f2a2cf0f Mon Sep 17 00:00:00 2001 From: "Paul S. Schweigert" Date: Fri, 24 Apr 2026 18:29:25 -0400 Subject: [PATCH 3/5] updates Signed-off-by: Paul S. Schweigert --- test/plugins/test_internal_tools.py | 54 ----------------------------- 1 file changed, 54 deletions(-) delete mode 100644 test/plugins/test_internal_tools.py diff --git a/test/plugins/test_internal_tools.py b/test/plugins/test_internal_tools.py deleted file mode 100644 index ec70bd79b..000000000 --- a/test/plugins/test_internal_tools.py +++ /dev/null @@ -1,54 +0,0 @@ -from typing import Any - -import pytest - -from mellea.backends.tools import MelleaTool -from mellea.core.backend import Backend -from mellea.core.base import GenerateLog, ModelOutputThunk, ModelToolCall -from mellea.plugins import HookType, PluginMode, block, hook, register -from mellea.stdlib.functional import _acall_tools - - -class _MockBackend(Backend): - async def _generate_from_context(self, *args, **kwargs): - pass - - async def generate_from_raw(self, *args, **kwargs): - pass - - -@pytest.mark.asyncio -async def test_tool_is_internal_accessible_in_hook(): - """Verify that is_internal flag is accessible in TOOL_PRE_INVOKE hook.""" - - observed_is_internal = [] - - @hook(HookType.TOOL_PRE_INVOKE) - async def checker(payload, _): - observed_is_internal.append(payload.model_tool_call.func.is_internal) - return None - - register(checker) - - # Create an internal tool - internal_tool = MelleaTool.from_callable( - lambda: "ok", name="internal", is_internal=True - ) - external_tool = MelleaTool.from_callable( - lambda: "ok", name="external", is_internal=False - ) - - tc_internal = ModelToolCall(name="internal", func=internal_tool, args={}) - tc_external = ModelToolCall(name="external", func=external_tool, args={}) - - mot = ModelOutputThunk( - value="", tool_calls={"internal": tc_internal, "external": tc_external} - ) - - await _acall_tools(mot, _MockBackend()) - - assert observed_is_internal == [True, False] - - -if __name__ == "__main__": - pytest.main([__file__]) From 1359bd93bc78de68efc9c34ef6ee0ea835d0d6a2 Mon Sep 17 00:00:00 2001 From: "Paul S. Schweigert" Date: Fri, 24 Apr 2026 18:29:47 -0400 Subject: [PATCH 4/5] updates Signed-off-by: Paul S. Schweigert --- docs/examples/plugins/tool_hooks.py | 5 ----- mellea/backends/tools.py | 14 +++----------- mellea/core/base.py | 4 ---- mellea/stdlib/components/react.py | 2 +- test/stdlib/frameworks/test_react_framework.py | 8 ++------ 5 files changed, 6 insertions(+), 27 deletions(-) diff --git a/docs/examples/plugins/tool_hooks.py b/docs/examples/plugins/tool_hooks.py index b1b4cd98e..fcd2e2fad 100644 --- a/docs/examples/plugins/tool_hooks.py +++ b/docs/examples/plugins/tool_hooks.py @@ -150,11 +150,6 @@ def parse_factor(): @hook(HookType.TOOL_PRE_INVOKE, mode=PluginMode.CONCURRENT, priority=5) async def enforce_tool_allowlist(payload, _): """Block any tool not on the explicit allow list.""" - # Internal tools (like ReAct's final_answer) are always permitted. - if payload.model_tool_call.func.is_internal: - log.info("[allowlist] permitted internal tool=%r", payload.model_tool_call.name) - return None - tool_name = payload.model_tool_call.name if tool_name not in ALLOWED_TOOLS: log.warning( diff --git a/mellea/backends/tools.py b/mellea/backends/tools.py index da1d57b16..66847f8fa 100644 --- a/mellea/backends/tools.py +++ b/mellea/backends/tools.py @@ -34,7 +34,6 @@ class MelleaTool(AbstractMelleaTool): tool_call (Callable): The underlying Python callable to invoke when the tool is run. as_json_tool (dict[str, Any]): The OpenAI-compatible JSON schema dict describing the tool's parameters. - is_internal (bool): Whether the tool is internal to Mellea. Defaults to ``False``. """ @@ -42,20 +41,14 @@ class MelleaTool(AbstractMelleaTool): # Our ModelToolCall is the class that has a reference to the tool and actually calls with arguments name: str - is_internal: bool _as_json_tool: dict[str, Any] _call_tool: Callable[..., Any] def __init__( - self, - name: str, - tool_call: Callable, - as_json_tool: dict[str, Any], - is_internal: bool = False, + self, name: str, tool_call: Callable, as_json_tool: dict[str, Any] ) -> None: """Initialize the tool with a name, tool call and as_json_tool dict.""" self.name = name - self.is_internal = is_internal self._as_json_tool = as_json_tool self._call_tool = tool_call @@ -180,7 +173,7 @@ def tool_call(*args, **kwargs): @classmethod def from_callable( - cls, func: Callable, name: str | None = None, is_internal: bool = False + cls, func: Callable, name: str | None = None ) -> "MelleaTool": """Create a MelleaTool from a plain Python callable. @@ -190,7 +183,6 @@ def from_callable( Args: func (Callable): The Python callable to wrap as a tool. name (str | None): Optional name override; defaults to ``func.__name__``. - is_internal (bool): Whether the tool is internal to Mellea. Defaults to ``False``. Returns: MelleaTool: A Mellea tool wrapping the callable. @@ -201,7 +193,7 @@ def from_callable( exclude_none=True ) tool_call = func - return MelleaTool(tool_name, tool_call, as_json, is_internal=is_internal) + return MelleaTool(tool_name, tool_call, as_json) @overload diff --git a/mellea/core/base.py b/mellea/core/base.py index 902861bf5..a3b40b97e 100644 --- a/mellea/core/base.py +++ b/mellea/core/base.py @@ -954,7 +954,6 @@ class AbstractMelleaTool(abc.ABC): Attributes: name (str): The unique name used to identify the tool in JSON descriptions and tool-call dispatch. - is_internal (bool): Whether the tool is internal to Mellea and should be bypassed by certain security policies. as_json_tool (dict[str, Any]): A JSON-serialisable description of the tool, compatible with the function-calling schemas expected by supported inference backends. """ @@ -962,9 +961,6 @@ class AbstractMelleaTool(abc.ABC): name: str """Name of the tool.""" - is_internal: bool = False - """Whether the tool is internal to Mellea.""" - @abc.abstractmethod def run(self, *args: Any, **kwargs: Any) -> Any: """Executes the tool with the provided arguments and returns the result. diff --git a/mellea/stdlib/components/react.py b/mellea/stdlib/components/react.py index 6d0234cfc..8f08fe8a0 100644 --- a/mellea/stdlib/components/react.py +++ b/mellea/stdlib/components/react.py @@ -71,7 +71,7 @@ def format_for_llm(self) -> TemplateRepresentation: ) finalizer_tool = MelleaTool.from_callable( - _mellea_finalize_tool, MELLEA_FINALIZER_TOOL, is_internal=True + _mellea_finalize_tool, MELLEA_FINALIZER_TOOL ) tools[MELLEA_FINALIZER_TOOL] = finalizer_tool diff --git a/test/stdlib/frameworks/test_react_framework.py b/test/stdlib/frameworks/test_react_framework.py index 5136772ba..e121a91f5 100644 --- a/test/stdlib/frameworks/test_react_framework.py +++ b/test/stdlib/frameworks/test_react_framework.py @@ -98,9 +98,7 @@ def _fn() -> str: def _final_answer_call(answer: str = "42") -> _ScriptedTurn: """Script a turn where the model calls final_answer with real arg flow.""" - tool = MelleaTool.from_callable( - _mellea_finalize_tool, MELLEA_FINALIZER_TOOL, is_internal=True - ) + tool = MelleaTool.from_callable(_mellea_finalize_tool, MELLEA_FINALIZER_TOOL) tc = ModelToolCall(name=MELLEA_FINALIZER_TOOL, func=tool, args={"answer": answer}) return _ScriptedTurn(value="", tool_calls={MELLEA_FINALIZER_TOOL: tc}) @@ -207,9 +205,7 @@ async def test_react_format_triggers_second_generation(): async def test_react_final_answer_with_extra_tool_rejected(): """final_answer alongside another tool in the same turn triggers assertion.""" search = _make_tool("search", "found it") - finalizer = MelleaTool.from_callable( - _mellea_finalize_tool, MELLEA_FINALIZER_TOOL, is_internal=True - ) + finalizer = MelleaTool.from_callable(_mellea_finalize_tool, MELLEA_FINALIZER_TOOL) both = { "search": ModelToolCall(name="search", func=search, args={}), MELLEA_FINALIZER_TOOL: ModelToolCall( From e628313c35276ade5d2c3ec319d2b6621de0f32e Mon Sep 17 00:00:00 2001 From: "Paul S. Schweigert" Date: Fri, 24 Apr 2026 18:40:35 -0400 Subject: [PATCH 5/5] lint Signed-off-by: Paul S. Schweigert --- mellea/backends/tools.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/mellea/backends/tools.py b/mellea/backends/tools.py index 66847f8fa..3dd056512 100644 --- a/mellea/backends/tools.py +++ b/mellea/backends/tools.py @@ -172,9 +172,7 @@ def tool_call(*args, **kwargs): ) from e @classmethod - def from_callable( - cls, func: Callable, name: str | None = None - ) -> "MelleaTool": + def from_callable(cls, func: Callable, name: str | None = None) -> "MelleaTool": """Create a MelleaTool from a plain Python callable. Introspects the callable's signature and docstring to build an