allow modeloptions.tools to be mapping or iterable; add func for extracting tools from list of actions

Author: jakelorocco

mellea/backends/tools.py

@@ -7,26 +7,58 @@
 from ollama._utils import convert_function_to_tool
 
 from mellea.backends.types import ModelOption
-from mellea.stdlib.base import Component, TemplateRepresentation
+from mellea.stdlib.base import CBlock, Component, TemplateRepresentation
 
 
 def add_tools_from_model_options(
     tools_dict: dict[str, Callable], model_options: dict[str, Any]
 ):
-    """If model_options has tools, it will add those tools to the tools_dict."""
+    """If model_options has tools, add those tools to the tools_dict."""
     model_opts_tools = model_options.get(ModelOption.TOOLS, None)
-
     if model_opts_tools is None:
         return
 
+    # Mappings are iterable.
     assert isinstance(model_opts_tools, Iterable), (
-        "ModelOption.TOOLS must be a list of Callables"
+        "ModelOption.TOOLS must be a list of Callables or dict[str, Callable]"
     )
-    for func in model_opts_tools:
-        assert callable(func), (
-            f"ModelOption.TOOLS must be a list of Callables, found {type(func)}"
-        )
-        tools_dict[func.__name__] = func
+
+    if isinstance(model_opts_tools, Mapping):
+        # Handle the dict case.
+        for func_name, func in model_opts_tools.items():
+            assert isinstance(func_name, str), (
+                f"If ModelOption.TOOLS is a dict, it must be a dict of [str, Callable]; found {type(func_name)} as the key instead"
+            )
+            assert callable(func), (
+                f"If ModelOption.TOOLS is a dict, it must be a dict of [str, Callable]; found {type(func)} as the value instead"
+            )
+            tools_dict[func_name] = func
+    else:
+        # Handle any other iterable / list here.
+        for func in model_opts_tools:
+            assert callable(func), (
+                f"If ModelOption.TOOLS is a list, it must be a list of Callables; found {type(func)}"
+            )
+            tools_dict[func.__name__] = func
+
+
+def add_tools_from_context_actions(
+    tools_dict: dict[str, Callable], ctx_actions: list[Component | CBlock] | None
+):
+    """If any of the actions in ctx_actions have tools in their template_representation, add those to the tools_dict."""
+    if ctx_actions is None:
+        return
+
+    for action in ctx_actions:
+        if not isinstance(action, Component):
+            continue  # Only components have template representations.
+
+        tr = action.format_for_llm()
+        if isinstance(tr, str) or tr.tools is None:
+            continue
+
+        for tool_name, func in tr.tools.items():
+            tools_dict[tool_name] = func
 
 
 def get_tools_from_action(action: Any) -> dict[str, Callable]:

mellea/backends/types.py

@@ -17,7 +17,7 @@ class ModelOption:
     """
 
     TOOLS = "@@@tools@@@"
-    """Must be a list of callables."""
+    """Must be a list of callables or a dict[str, Callable]."""
 
     MAX_NEW_TOKENS = "@@@max_new_tokens@@@"
     SYSTEM_PROMPT = "@@@system_prompt@@@"

test/test_tool_calls.py

@@ -2,9 +2,9 @@
 
 from mellea.backends import Backend
 from mellea.backends.ollama import OllamaModelBackend
-from mellea.backends.tools import add_tools_from_model_options
+from mellea.backends.tools import add_tools_from_context_actions, add_tools_from_model_options
 from mellea.backends.types import ModelOption
-from mellea.stdlib.base import ModelOutputThunk
+from mellea.stdlib.base import CBlock, Component, ModelOutputThunk, TemplateRepresentation
 from mellea.stdlib.docs.richdocument import Table
 from mellea.stdlib.session import LinearContext, MelleaSession
 
@@ -29,13 +29,47 @@ def table() -> Table:
     assert t is not None, "test setup failed: could not create table from markdown"
     return t
 
-def test_add_tools_from_model_options(table: Table):
+class FakeToolComponent(Component):
+    def __init__(self) -> None:
+        super().__init__()
+    
+    def tool1(self):
+        return
+    
+    def parts(self):
+        return []
+
+    def format_for_llm(self) -> TemplateRepresentation:
+        return TemplateRepresentation(
+            obj=self,
+            args={"arg": None},
+            tools={
+                self.tool1.__name__: self.tool1
+            }
+        )
+
+class FakeToolComponentWithExtraTool(FakeToolComponent):
+    def __init__(self) -> None:
+        super().__init__()
+    
+    def tool2(self):
+        return
+    
+    def format_for_llm(self) -> TemplateRepresentation:
+        tr = super().format_for_llm()
+        assert tr.tools is not None
+        tr.tools[self.tool2.__name__] = self.tool2
+        return tr
+
+
+def test_add_tools_from_model_options_list():
     def get_weather(location: str) -> int:
         """Returns the weather in Celsius."""
         return 21
 
+    ftc = FakeToolComponent()
     model_options = {
-        ModelOption.TOOLS: [get_weather, table.content_as_string]
+        ModelOption.TOOLS: [get_weather, ftc.tool1]
     }
 
     tools = {}
@@ -44,36 +78,75 @@ def get_weather(location: str) -> int:
     assert tools["get_weather"] == get_weather
 
     # Must use `==` for bound methods.
-    assert tools["content_as_string"] == table.content_as_string, f"{tools["content_as_string"]} is not {table.content_as_string}"
+    assert tools["tool1"] == ftc.tool1, f"{tools["tool1"]} should == {ftc.tool1}"
 
-def test_tool_called(m: MelleaSession, table: Table):
-    """We don't force tools to be called. As a result, this test might unexpectedly fail."""
-    r = 10
 
-    returned_tool = False
-    for i in range(r):
-        transformed = m.transform(table, "add a new row to this table")
-        if isinstance(transformed, Table):
-            returned_tool = True
-            break
+def test_add_tools_from_model_options_map():
+    def get_weather(location: str) -> int:
+        """Returns the weather in Celsius."""
+        return 21
+
+    ftc = FakeToolComponent()
+    model_options = {
+        ModelOption.TOOLS: {
+            get_weather.__name__: get_weather,
+            ftc.tool1.__name__: ftc.tool1
+        }
+    }
+
+    tools = {}
+    add_tools_from_model_options(tools, model_options)
+
+    assert tools["get_weather"] == get_weather
+
+    # Must use `==` for bound methods.
+    assert tools["tool1"] == ftc.tool1, f"{tools["tool1"]} should == {ftc.tool1}"
+
+
+def test_add_tools_from_context_actions():
+
+    ftc1 = FakeToolComponentWithExtraTool()
+    ftc2 = FakeToolComponent()
+
+    ctx_actions = [CBlock("Hello"), ftc1, ftc2]
+    tools = {}
+    add_tools_from_context_actions(tools, ctx_actions)
+
+    # Check that tools with the same name get properly overwritten in order of ctx.
+    assert tools["tool1"] == ftc2.tool1, f"{tools["tool1"]} should == {ftc2.tool1}"
+
+    # Check that tools that aren't overwritten are still there.
+    assert tools["tool2"] == ftc1.tool2, f"{tools["tool2"]} should == {ftc1.tool2}"
+
+
+# def test_tool_called(m: MelleaSession, table: Table):
+#     """We don't force tools to be called. As a result, this test might unexpectedly fail."""
+#     r = 10
+
+#     returned_tool = False
+#     for i in range(r):
+#         transformed = m.transform(table, "add a new row to this table")
+#         if isinstance(transformed, Table):
+#             returned_tool = True
+#             break
 
-    assert returned_tool, f"did not return a tool after {r} attempts"
+#     assert returned_tool, f"did not return a tool after {r} attempts"
 
 
-def test_tool_not_called(m: MelleaSession, table: Table):
-    """Ensure tools aren't always called when provided."""
-    r = 10
+# def test_tool_not_called(m: MelleaSession, table: Table):
+#     """Ensure tools aren't always called when provided."""
+#     r = 10
 
-    returned_no_tool = False
-    for i in range(r):
-        transformed = m.transform(table, "output a text description of this table")
-        if isinstance(transformed, ModelOutputThunk):
-            returned_no_tool = True
-            break
+#     returned_no_tool = False
+#     for i in range(r):
+#         transformed = m.transform(table, "output a text description of this table")
+#         if isinstance(transformed, ModelOutputThunk):
+#             returned_no_tool = True
+#             break
 
-    assert (
-        returned_no_tool
-    ), f"only returned tools after {r} attempts, should've returned a response with no tools"
+#     assert (
+#         returned_no_tool
+#     ), f"only returned tools after {r} attempts, should've returned a response with no tools"
 
 if __name__ == "__main__":
     pytest.main([__file__])