update handling of tools for the given action

Author: jakelorocco

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
@@ -35,7 +34,6 @@
     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
@@ -47,7 +45,6 @@
     GenerateLog,
     ModelOutputThunk,
     ModelToolCall,
-    TemplateRepresentation,
 )
 from mellea.stdlib.chat import Message
 from mellea.stdlib.requirement import ALoraRequirement, LLMaJRequirement, Requirement
@@ -327,13 +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)
-                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()}")
+                    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
 
@@ -16,7 +15,6 @@
 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
@@ -299,13 +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)
-
                 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

@@ -24,7 +24,6 @@
     add_tools_from_context_actions,
     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
@@ -35,7 +34,6 @@
     GenerateLog,
     ModelOutputThunk,
     ModelToolCall,
-    TemplateRepresentation,
 )
 from mellea.stdlib.chat import Message
 from mellea.stdlib.requirement import ALoraRequirement, LLMaJRequirement, Requirement
@@ -409,14 +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)
-
                 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()}")
+
+                # 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

@@ -54,28 +54,13 @@ def add_tools_from_context_actions(
             continue  # Only components have template representations.
 
         tr = action.format_for_llm()
-        if isinstance(tr, str) or tr.tools is None:
+        if not isinstance(tr, TemplateRepresentation) 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]:
-    """If an object is a Component with a TemplateRepresentation, grabs it's tools field.
-
-    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
-
-    return {}
-
-
 def convert_tools_to_json(tools: dict[str, Callable]) -> list[dict]:
     """Convert tools to json dict representation.
 

mellea/backends/watsonx.py

@@ -17,7 +17,6 @@
     add_tools_from_context_actions,
     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
@@ -266,9 +265,12 @@ 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)
                 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)