switch ModelOptions.TOOLS to be a list; standardize tool handling; add tests

Author: jakelorocco

mellea/backends/huggingface.py

@@ -32,6 +32,7 @@
 from mellea.backends.formatter import Formatter, FormatterBackend, TemplateFormatter
 from mellea.backends.model_ids import ModelIdentifier
 from mellea.backends.tools import (
+    add_tools_from_model_options,
     convert_tools_to_json,
     get_tools_from_action,
     parse_tools,
@@ -329,24 +330,8 @@ def _generate_from_context_standard(
                         action.format_for_llm(), TemplateRepresentation
                     ):
                         tools = get_tools_from_action(action)
-
-                    model_options_tools = model_options.get(ModelOption.TOOLS, None)
-                    if model_options_tools is not None:
-                        assert isinstance(model_options_tools, dict)
-                        for fn_name in model_options_tools:
-                            # invariant re: relationship between the model_options set of tools and the TemplateRepresentation set of tools
-                            assert fn_name not in tools.keys(), (
-                                f"Cannot add tool {fn_name} because that tool was already defined in the TemplateRepresentation for the action."
-                            )
-                            # type checking because ModelOptions is an untyped dict and the calling convention for tools isn't clearly documented at our abstraction boundaries.
-                            assert type(fn_name) is str, (
-                                "When providing a `ModelOption.TOOLS` parameter to `model_options`, always used the type Dict[str, Callable] where `str` is the function name and the callable is the function."
-                            )
-                            assert callable(model_options_tools[fn_name]), (
-                                "When providing a `ModelOption.TOOLS` parameter to `model_options`, always used the type Dict[str, Callable] where `str` is the function name and the callable is the function."
-                            )
-                            # Add the model_options tool to the existing set of tools.
-                            tools[fn_name] = model_options_tools[fn_name]
+                add_tools_from_model_options(tools, model_options)
+            FancyLogger.get_logger().info(f"Tools for call: {tools.keys()}")
 
             seed = model_options.get(ModelOption.SEED, None)
             if seed is not None:

mellea/backends/ollama.py

@@ -13,7 +13,7 @@
 from mellea.backends import BaseModelSubclass
 from mellea.backends.formatter import Formatter, FormatterBackend, TemplateFormatter
 from mellea.backends.model_ids import ModelIdentifier
-from mellea.backends.tools import get_tools_from_action
+from mellea.backends.tools import add_tools_from_model_options, get_tools_from_action
 from mellea.backends.types import ModelOption
 from mellea.helpers.fancy_logger import FancyLogger
 from mellea.stdlib.base import (
@@ -300,23 +300,7 @@ def generate_from_chat_context(
                 ):
                     tools = get_tools_from_action(action)
 
-                model_options_tools = model_opts.get(ModelOption.TOOLS, None)
-                if model_options_tools is not None:
-                    assert isinstance(model_options_tools, dict)
-                    for fn_name in model_options_tools:
-                        # invariant re: relationship between the model_options set of tools and the TemplateRepresentation set of tools
-                        assert fn_name not in tools.keys(), (
-                            f"Cannot add tool {fn_name} because that tool was already defined in the TemplateRepresentation for the action."
-                        )
-                        # type checking because ModelOptions is an untyped dict and the calling convention for tools isn't clearly documented at our abstraction boundaries.
-                        assert type(fn_name) is str, (
-                            "When providing a `ModelOption.TOOLS` parameter to `model_options`, always used the type Dict[str, Callable] where `str` is the function name and the callable is the function."
-                        )
-                        assert callable(model_options_tools[fn_name]), (
-                            "When providing a `ModelOption.TOOLS` parameter to `model_options`, always used the type Dict[str, Callable] where `str` is the function name and the callable is the function."
-                        )
-                        # Add the model_options tool to the existing set of tools.
-                        tools[fn_name] = model_options_tools[fn_name]
+                add_tools_from_model_options(tools, model_opts)
             FancyLogger.get_logger().info(f"Tools for call: {tools.keys()}")
 
         # Generate a chat response from ollama, using the chat messages.

mellea/backends/openai.py

@@ -20,7 +20,11 @@
 from mellea.backends.aloras import Alora, AloraBackendMixin
 from mellea.backends.formatter import Formatter, FormatterBackend, TemplateFormatter
 from mellea.backends.model_ids import ModelIdentifier
-from mellea.backends.tools import convert_tools_to_json, get_tools_from_action
+from mellea.backends.tools import (
+    add_tools_from_model_options,
+    convert_tools_to_json,
+    get_tools_from_action,
+)
 from mellea.backends.types import ModelOption
 from mellea.helpers.fancy_logger import FancyLogger
 from mellea.stdlib.base import (
@@ -409,23 +413,8 @@ def _generate_from_chat_context_standard(
                 ):
                     tools = get_tools_from_action(action)
 
-                model_options_tools = model_opts.get(ModelOption.TOOLS, None)
-                if model_options_tools is not None:
-                    assert isinstance(model_options_tools, dict)
-                    for fn_name in model_options_tools:
-                        # invariant re: relationship between the model_options set of tools and the TemplateRepresentation set of tools
-                        assert fn_name not in tools.keys(), (
-                            f"Cannot add tool {fn_name} because that tool was already defined in the TemplateRepresentation for the action."
-                        )
-                        # type checking because ModelOptions is an untyped dict and the calling convention for tools isn't clearly documented at our abstraction boundaries.
-                        assert type(fn_name) is str, (
-                            "When providing a `ModelOption.TOOLS` parameter to `model_options`, always used the type Dict[str, Callable] where `str` is the function name and the callable is the function."
-                        )
-                        assert callable(model_options_tools[fn_name]), (
-                            "When providing a `ModelOption.TOOLS` parameter to `model_options`, always used the type Dict[str, Callable] where `str` is the function name and the callable is the function."
-                        )
-                        # Add the model_options tool to the existing set of tools.
-                        tools[fn_name] = model_options_tools[fn_name]
+                add_tools_from_model_options(tools, model_opts)
+        FancyLogger.get_logger().info(f"Tools for call: {tools.keys()}")
 
         thinking = model_opts.get(ModelOption.THINKING, None)
         if type(thinking) is bool and thinking:

mellea/backends/tools.py

@@ -1,14 +1,34 @@
 """Utilities for dealing with tools."""
 
 import json
-from collections.abc import Callable, Generator, Mapping
+from collections.abc import Callable, Generator, Iterable, Mapping
 from typing import Any
 
 from ollama._utils import convert_function_to_tool
 
+from mellea.backends.types import ModelOption
 from mellea.stdlib.base import 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."""
+    model_opts_tools = model_options.get(ModelOption.TOOLS, None)
+
+    if model_opts_tools is None:
+        return
+
+    assert isinstance(model_opts_tools, Iterable), (
+        "ModelOption.TOOLS must be a list of Callables"
+    )
+    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
+
+
 def get_tools_from_action(action: Any) -> dict[str, Callable]:
     """If an object is a Component with a TemplateRepresentation, grabs it's tools field.
 

mellea/backends/types.py

@@ -17,6 +17,8 @@ class ModelOption:
     """
 
     TOOLS = "@@@tools@@@"
+    """Must be a list of callables."""
+
     MAX_NEW_TOKENS = "@@@max_new_tokens@@@"
     SYSTEM_PROMPT = "@@@system_prompt@@@"
     TEMPERATURE = "temperature"

mellea/backends/watsonx.py

@@ -13,7 +13,11 @@
 from mellea.backends import BaseModelSubclass, model_ids
 from mellea.backends.formatter import Formatter, FormatterBackend, TemplateFormatter
 from mellea.backends.model_ids import ModelIdentifier
-from mellea.backends.tools import convert_tools_to_json, get_tools_from_action
+from mellea.backends.tools import (
+    add_tools_from_model_options,
+    convert_tools_to_json,
+    get_tools_from_action,
+)
 from mellea.backends.types import ModelOption
 from mellea.helpers.fancy_logger import FancyLogger
 from mellea.stdlib.base import (
@@ -262,24 +266,9 @@ def generate_from_chat_context(
                 )
             else:
                 tools = get_tools_from_action(action)
+                add_tools_from_model_options(tools, model_opts)
 
-                model_options_tools = model_opts.get(ModelOption.TOOLS, None)
-                if model_options_tools is not None:
-                    assert isinstance(model_options_tools, dict)
-                    for fn_name in model_options_tools:
-                        # invariant re: relationship between the model_options set of tools and the TemplateRepresentation set of tools
-                        assert fn_name not in tools.keys(), (
-                            f"Cannot add tool {fn_name} because that tool was already defined in the TemplateRepresentation for the action."
-                        )
-                        # type checking because ModelOptions is an untyped dict and the calling convention for tools isn't clearly documented at our abstraction boundaries.
-                        assert type(fn_name) is str, (
-                            "When providing a `ModelOption.TOOLS` parameter to `model_options`, always used the type Dict[str, Callable] where `str` is the function name and the callable is the function."
-                        )
-                        assert callable(model_options_tools[fn_name]), (
-                            "When providing a `ModelOption.TOOLS` parameter to `model_options`, always used the type Dict[str, Callable] where `str` is the function name and the callable is the function."
-                        )
-                        # Add the model_options tool to the existing set of tools.
-                        tools[fn_name] = model_options_tools[fn_name]
+            FancyLogger.get_logger().info(f"Tools for call: {tools.keys()}")
 
         formatted_tools = convert_tools_to_json(tools)
         chat_response = self._model.chat(

test/test_tool_calls.py

@@ -2,6 +2,8 @@
 
 from mellea.backends import Backend
 from mellea.backends.ollama import OllamaModelBackend
+from mellea.backends.tools import add_tools_from_model_options
+from mellea.backends.types import ModelOption
 from mellea.stdlib.base import ModelOutputThunk
 from mellea.stdlib.docs.richdocument import Table
 from mellea.stdlib.session import LinearContext, MelleaSession
@@ -27,6 +29,22 @@ 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):
+    def get_weather(location: str) -> int:
+        """Returns the weather in Celsius."""
+        return 21
+
+    model_options = {
+        ModelOption.TOOLS: [get_weather, table.content_as_string]
+    }
+
+    tools = {}
+    add_tools_from_model_options(tools, model_options)
+
+    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}"
 
 def test_tool_called(m: MelleaSession, table: Table):
     """We don't force tools to be called. As a result, this test might unexpectedly fail."""