overhaul tools for model calls

docs/tutorial.md

@@ -20,6 +20,7 @@
 - [Chapter 9: Interoperability with Other Frameworks](#chapter-9-interoperability-with-other-frameworks)
 - [Chapter 10: Prompt Engineering for Mellea](#chapter-10-prompt-engineering-for-m)
   - [Custom  Templates](#custom-templates)
+- [Chapter 11: Tool Calling](#chapter-11-tool-calling)
 - [Appendix: Contributing to Melles](#appendix-contributing-to-mellea)
 
 ## Chapter 1: What Is Generative Programming
@@ -1281,6 +1282,33 @@ To customize the template and template representation of an existing class, simp
 
 See [`mellea/docs/examples/mify/rich_document_advanced.py`](./examples/mify/rich_document_advanced.py)
 
+## Chapter 11: Tool Calling
+Mellea supports tool calling for providers/models that support it. Most session level functions support setting a tool_calls boolean. Setting this to true allows tools to be called, but there's no guarantee that a model will call them.
+
+Tools can be made available for the model to call in a few ways:
+1. Components: components can have a TemplateRepresentation object that contains tools.
+2. Context: depending on the context, the components in that context can be used as sources of additional tools in the exact same way they would if they were the current action.
+3. `ModelOptions.TOOLS`: model options can include a tools parameter. The preferred way of passing these tools is as a list of function objects.
+
+Currently, tools are identified by the name of the function. If there are conflicts, the most recent tool with that name will be preferred. This means the tools available to the model will have the same priority listed above:
+1. Tools from the current component will always be included
+2. Tools from the context will be included if there are no name conflicts. A given context can decide what tools to surface, but in most cases, tools from the most recent component in the context will take priority over tools from older requests.
+3. Tools from `ModelOptions.TOOLS` will only be added if they do not conflict with any of the above functions.
+
+For examples on adding tools to the template representation of a component, see the `Table` object in [richdocument.py](../mellea/stdlib/docs/richdocument.py).
+
+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
+from mellea.backends.types import ModelOption
+
+def web_search(query: str) -> str:
+    ...
+
+model_opts = {
+    ModelOptions.TOOLS: [web_search]
+}
+```
+
 ## Appendix: Contributing to Mellea
 
 ### Contributor Guide: Requirements and Verifiers

mellea/backends/huggingface.py

@@ -10,9 +10,8 @@
 import datetime
 import inspect
 import json
-import os
 from collections.abc import Callable
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any
 
 import outlines
 import outlines_core
@@ -32,8 +31,9 @@
 from mellea.backends.formatter import Formatter, FormatterBackend, TemplateFormatter
 from mellea.backends.model_ids import ModelIdentifier
 from mellea.backends.tools import (
+    add_tools_from_context_actions,
+    add_tools_from_model_options,
     convert_tools_to_json,
-    get_tools_from_action,
     parse_tools,
 )
 from mellea.backends.types import ModelOption
@@ -45,7 +45,6 @@
     GenerateLog,
     ModelOutputThunk,
     ModelToolCall,
-    TemplateRepresentation,
 )
 from mellea.stdlib.chat import Message
 from mellea.stdlib.requirement import ALoraRequirement, LLMaJRequirement, Requirement
@@ -325,28 +324,15 @@ def _generate_from_context_standard(
                         f"Tool calling typically uses constrained generation, but you have specified a `format` in your generate call. NB: tool calling is superseded by format; we will NOT call tools for your request: {action}"
                     )
                 else:
-                    if isinstance(action, Component) and isinstance(
-                        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)
+                    add_tools_from_context_actions(
+                        tools, ctx.actions_for_available_tools()
+                    )
+
+                    # Add the tools from the action for this generation last so that
+                    # they overwrite conflicting names.
+                    add_tools_from_context_actions(tools, [action])
+                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

@@ -2,7 +2,6 @@
 
 import asyncio
 import datetime
-import os
 from collections.abc import Callable
 from typing import Any
 
@@ -13,7 +12,10 @@
 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_context_actions,
+    add_tools_from_model_options,
+)
 from mellea.backends.types import ModelOption
 from mellea.helpers.fancy_logger import FancyLogger
 from mellea.stdlib.base import (
@@ -295,28 +297,12 @@ def generate_from_chat_context(
                     f"Tool calling typically uses constrained generation, but you have specified a `format` in your generate call. NB: tool calling is superseded by format; we will NOT call tools for your request: {action}"
                 )
             else:
-                if isinstance(action, Component) and isinstance(
-                    action.format_for_llm(), TemplateRepresentation
-                ):
-                    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)
+                add_tools_from_context_actions(tools, ctx.actions_for_available_tools())
+
+                # Add the tools from the action for this generation last so that
+                # they overwrite conflicting names.
+                add_tools_from_context_actions(tools, [action])
             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_context_actions,
+    add_tools_from_model_options,
+    convert_tools_to_json,
+)
 from mellea.backends.types import ModelOption
 from mellea.helpers.fancy_logger import FancyLogger
 from mellea.stdlib.base import (
@@ -30,7 +34,6 @@
     GenerateLog,
     ModelOutputThunk,
     ModelToolCall,
-    TemplateRepresentation,
 )
 from mellea.stdlib.chat import Message
 from mellea.stdlib.requirement import ALoraRequirement, LLMaJRequirement, Requirement
@@ -404,28 +407,13 @@ def _generate_from_chat_context_standard(
                     f"Tool calling typically uses constrained generation, but you have specified a `format` in your generate call. NB: tool calling is superseded by format; we will NOT call tools for your request: {action}"
                 )
             else:
-                if isinstance(action, Component) and isinstance(
-                    action.format_for_llm(), TemplateRepresentation
-                ):
-                    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)
+                add_tools_from_context_actions(tools, ctx.actions_for_available_tools())
+
+                # Add the tools from the action for this generation last so that
+                # they overwrite conflicting names.
+                add_tools_from_context_actions(tools, [action])
+            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,27 +1,64 @@
 """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.stdlib.base import Component, TemplateRepresentation
+from mellea.backends.types import ModelOption
+from mellea.stdlib.base import CBlock, Component, TemplateRepresentation
 
 
-def get_tools_from_action(action: Any) -> dict[str, Callable]:
-    """If an object is a Component with a TemplateRepresentation, grabs it's tools field.
+def add_tools_from_model_options(
+    tools_dict: dict[str, Callable], model_options: dict[str, Any]
+):
+    """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 or dict[str, Callable]"
+    )
+
+    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.
 
-    Returns:
-        dict: mapping function names to callables
-    """
-    if isinstance(action, Component):
         tr = action.format_for_llm()
-        if isinstance(tr, TemplateRepresentation):
-            if tr.tools:
-                return tr.tools
+        if not isinstance(tr, TemplateRepresentation) or tr.tools is None:
+            continue
 
-    return {}
+        for tool_name, func in tr.tools.items():
+            tools_dict[tool_name] = func
 
 
 def convert_tools_to_json(tools: dict[str, Callable]) -> list[dict]:

mellea/backends/types.py

@@ -17,6 +17,8 @@ class ModelOption:
     """
 
     TOOLS = "@@@tools@@@"
+    """Must be a list of callables or a dict[str, Callable]."""
+
     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_context_actions,
+    add_tools_from_model_options,
+    convert_tools_to_json,
+)
 from mellea.backends.types import ModelOption
 from mellea.helpers.fancy_logger import FancyLogger
 from mellea.stdlib.base import (
@@ -261,25 +265,13 @@ def generate_from_chat_context(
                     f"tool calling is superseded by format; will not call tools for request: {action}"
                 )
             else:
-                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)
+                add_tools_from_context_actions(tools, ctx.actions_for_available_tools())
+
+                # Add the tools from the action for this generation last so that
+                # they overwrite conflicting names.
+                add_tools_from_context_actions(tools, [action])
+            FancyLogger.get_logger().info(f"Tools for call: {tools.keys()}")
 
         formatted_tools = convert_tools_to_json(tools)
         chat_response = self._model.chat(