add tools from context to backend calls; add test for tool from context calling

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_context_actions,
     add_tools_from_model_options,
     convert_tools_to_json,
     get_tools_from_action,
@@ -331,6 +332,7 @@ def _generate_from_context_standard(
                     ):
                         tools = get_tools_from_action(action)
                 add_tools_from_model_options(tools, model_options)
+                add_tools_from_context_actions(tools, ctx.actions_for_available_tools())
             FancyLogger.get_logger().info(f"Tools for call: {tools.keys()}")
 
             seed = model_options.get(ModelOption.SEED, None)

mellea/backends/ollama.py

@@ -13,7 +13,11 @@
 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 add_tools_from_model_options, get_tools_from_action
+from mellea.backends.tools import (
+    add_tools_from_context_actions,
+    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 (
@@ -301,6 +305,7 @@ def generate_from_chat_context(
                     tools = get_tools_from_action(action)
 
                 add_tools_from_model_options(tools, model_opts)
+                add_tools_from_context_actions(tools, ctx.actions_for_available_tools())
             FancyLogger.get_logger().info(f"Tools for call: {tools.keys()}")
 
         # Generate a chat response from ollama, using the chat messages.

mellea/backends/openai.py

@@ -21,6 +21,7 @@
 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,
@@ -414,6 +415,7 @@ def _generate_from_chat_context_standard(
                     tools = get_tools_from_action(action)
 
                 add_tools_from_model_options(tools, model_opts)
+                add_tools_from_context_actions(tools, ctx.actions_for_available_tools())
         FancyLogger.get_logger().info(f"Tools for call: {tools.keys()}")
 
         thinking = model_opts.get(ModelOption.THINKING, None)

mellea/backends/watsonx.py

@@ -14,6 +14,7 @@
 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,
@@ -267,7 +268,7 @@ def generate_from_chat_context(
             else:
                 tools = get_tools_from_action(action)
                 add_tools_from_model_options(tools, model_opts)
-
+                add_tools_from_context_actions(tools, ctx.actions_for_available_tools())
             FancyLogger.get_logger().info(f"Tools for call: {tools.keys()}")
 
         formatted_tools = convert_tools_to_json(tools)

test/test_tool_calls.py

@@ -119,9 +119,34 @@ def test_add_tools_from_context_actions():
     assert tools["tool2"] == ftc1.tool2, f"{tools["tool2"]} should == {ftc1.tool2}"
 
 
+def test_tool_called_from_context_action(m: MelleaSession, table: Table):
+    """Make sure tools can be called from actions in the context."""
+    r = 10
+    m.ctx.reset()
+
+    # Insert a component with tools into the context.
+    m.ctx.insert(table)
+
+    returned_tool = False
+    for i in range(r):
+        # Make sure the specific generate call is on a different action with
+        # no tools to make sure it's a tool from the context.
+        result = m.backend.generate_from_context(
+            CBlock("Add a row to the table."),
+            m.ctx,
+            tool_calls=True
+        )
+        if result.tool_calls is not None and len(result.tool_calls) > 0:
+            returned_tool = True
+            break
+
+    assert returned_tool, f"did not return a tool after {r} attempts"
+
+
 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
+    m.ctx.reset()
 
     returned_tool = False
     for i in range(r):
@@ -136,6 +161,7 @@ def test_tool_called(m: MelleaSession, table: Table):
 def test_tool_not_called(m: MelleaSession, table: Table):
     """Ensure tools aren't always called when provided."""
     r = 10
+    m.ctx.reset()
 
     returned_no_tool = False
     for i in range(r):