feat: new MelleaTool class and adoption across mellea

mellea/backends/huggingface.py

@@ -40,6 +40,7 @@
     ModelOutputThunk,
     Requirement,
 )
+from ..core.base import AbstractMelleaTool
 from ..formatters import ChatFormatter, TemplateFormatter
 from ..helpers import message_to_openai_message, messages_to_docs, send_to_queue
 from ..stdlib.components import Intrinsic, Message
@@ -564,7 +565,7 @@ async def _generate_from_context_with_kv_cache(
             ctx_as_chat = to_chat(action, ctx, self.formatter, system_prompt)
 
             # Append tool call information if applicable.
-            tools: dict[str, Callable] = dict()
+            tools: dict[str, AbstractMelleaTool] = dict()
             if tool_calls:
                 if _format:
                     FancyLogger.get_logger().warning(
@@ -723,7 +724,7 @@ async def _generate_from_context_standard(
             ctx_as_chat = to_chat(action, ctx, self.formatter, system_prompt)
 
             # Append tool call information if applicable.
-            tools: dict[str, Callable] = dict()
+            tools: dict[str, AbstractMelleaTool] = dict()
             if tool_calls:
                 if _format:
                     FancyLogger.get_logger().warning(
@@ -883,7 +884,7 @@ async def post_processing(
         conversation: list[dict],
         _format: type[BaseModelSubclass] | None,
         tool_calls: bool,
-        tools: dict[str, Callable],
+        tools: dict[str, AbstractMelleaTool],
         seed,
         input_ids,
     ):

mellea/backends/litellm.py

@@ -24,6 +24,7 @@
     ModelOutputThunk,
     ModelToolCall,
 )
+from ..core.base import AbstractMelleaTool
 from ..formatters import ChatFormatter, TemplateFormatter
 from ..helpers import (
     chat_completion_delta_merge,
@@ -400,7 +401,7 @@ async def post_processing(
         self,
         mot: ModelOutputThunk,
         conversation: list[dict],
-        tools: dict[str, Callable],
+        tools: dict[str, AbstractMelleaTool],
         thinking,
         _format,
     ):
@@ -457,8 +458,8 @@ async def post_processing(
     @staticmethod
     def _extract_tools(
         action, _format, model_opts, tool_calls, ctx
-    ) -> dict[str, Callable]:
-        tools: dict[str, Callable] = dict()
+    ) -> dict[str, AbstractMelleaTool]:
+        tools: dict[str, AbstractMelleaTool] = dict()
         if tool_calls:
             if _format:
                 FancyLogger.get_logger().warning(
@@ -579,7 +580,7 @@ async def generate_from_raw(
 
     def _extract_model_tool_requests(
         self,
-        tools: dict[str, Callable],
+        tools: dict[str, AbstractMelleaTool],
         chat_response: litellm.ModelResponse,  # type: ignore
     ) -> dict[str, ModelToolCall] | None:
         model_tool_calls: dict[str, ModelToolCall] = {}

mellea/backends/ollama.py

@@ -22,6 +22,7 @@
     ModelOutputThunk,
     ModelToolCall,
 )
+from ..core.base import AbstractMelleaTool
 from ..formatters import ChatFormatter, TemplateFormatter
 from ..helpers import ClientCache, get_current_event_loop, send_to_queue
 from ..stdlib.components import Message
@@ -322,7 +323,7 @@ async def generate_from_chat_context(
         )
 
         # Append tool call information if applicable.
-        tools: dict[str, Callable] = dict()
+        tools: dict[str, AbstractMelleaTool] = dict()
         if tool_calls:
             if _format:
                 FancyLogger.get_logger().warning(
@@ -343,7 +344,7 @@ async def generate_from_chat_context(
         ] = self._async_client.chat(
             model=self._get_ollama_model_id(),
             messages=conversation,
-            tools=list(tools.values()),
+            tools=[t.as_json_tool for t in tools.values()],
             think=model_opts.get(ModelOption.THINKING, None),
             stream=model_opts.get(ModelOption.STREAM, False),
             options=self._make_backend_specific_and_remove(model_opts),
@@ -499,7 +500,7 @@ async def generate_from_raw(
         return results
 
     def _extract_model_tool_requests(
-        self, tools: dict[str, Callable], chat_response: ollama.ChatResponse
+        self, tools: dict[str, AbstractMelleaTool], chat_response: ollama.ChatResponse
     ) -> dict[str, ModelToolCall] | None:
         model_tool_calls: dict[str, ModelToolCall] = {}
 
@@ -525,7 +526,7 @@ async def processing(
         self,
         mot: ModelOutputThunk,
         chunk: ollama.ChatResponse,
-        tools: dict[str, Callable],
+        tools: dict[str, AbstractMelleaTool],
     ):
         """Called during generation to add information from a single ChatResponse to the ModelOutputThunk."""
         if mot._thinking is None:
@@ -557,7 +558,7 @@ async def post_processing(
         self,
         mot: ModelOutputThunk,
         conversation: list[dict],
-        tools: dict[str, Callable],
+        tools: dict[str, AbstractMelleaTool],
         _format,
     ):
         """Called when generation is done."""

mellea/backends/openai.py

@@ -28,6 +28,7 @@
     ModelOutputThunk,
     Requirement,
 )
+from ..core.base import AbstractMelleaTool
 from ..formatters import ChatFormatter, TemplateFormatter
 from ..helpers import (
     ClientCache,
@@ -591,7 +592,7 @@ async def _generate_from_chat_context_standard(
                 }
 
         # Append tool call information if applicable.
-        tools: dict[str, Callable] = dict()
+        tools: dict[str, AbstractMelleaTool] = dict()
         if tool_calls:
             if _format:
                 FancyLogger.get_logger().warning(
@@ -713,7 +714,7 @@ async def processing(
     async def post_processing(
         self,
         mot: ModelOutputThunk,
-        tools: dict[str, Callable],
+        tools: dict[str, AbstractMelleaTool],
         conversation: list[dict],
         thinking,
         seed,

mellea/backends/tools.py

@@ -1,4 +1,4 @@
-"""Utilities for dealing with tools."""
+"""Utilities for dealing with LLM tools."""
 
 import inspect
 import json
@@ -9,12 +9,88 @@
 
 from pydantic import BaseModel, ConfigDict, Field
 
+from mellea.core.utils import FancyLogger
+
 from ..core import CBlock, Component, TemplateRepresentation
+from ..core.base import AbstractMelleaTool
 from .model_options import ModelOption
 
 
+class MelleaTool(AbstractMelleaTool):
+    """Tool class to represent a tool."""
+
+    # Tool is what we pass as a model option / as input
+    # Our ModelToolCall is the class that has a reference to the tool and actually calls with arguments
+
+    name: str
+    _as_json_tool: dict[str, Any]
+    _call_tool: Callable[..., Any]
+
+    def __init__(
+        self, name: str, tool_call: Callable, as_json_tool: dict[str, Any]
+    ) -> None:
+        """Initialize the tool with a name, tool call and as_json_tool dict."""
+        self.name = name
+        self._as_json_tool = as_json_tool
+        self._call_tool = tool_call
+
+    def run(self, *args, **kwargs) -> Any:
+        """Run the tool with the given arguments."""
+        return self._call_tool(*args, **kwargs)
+
+    @property
+    def as_json_tool(self) -> dict[str, Any]:
+        """Return the tool converted to a OpenAI compatible JSON object."""
+        return self._as_json_tool
+
+    @classmethod
+    def from_langchain(cls, tool: Any):
+        """Create a Tool from a langchain tool object."""
+        try:
+            from langchain_core.tools import BaseTool
+            from langchain_core.utils.function_calling import convert_to_openai_tool
+
+            if isinstance(tool, BaseTool):
+                tool_name = tool.name
+                as_json = convert_to_openai_tool(tool)
+
+                def parameter_remapper(*args, **kwargs):
+                    """Langchain tools expect their first argument to be 'tool_input'."""
+                    if args is not None or len(args) != 0:
+                        # This shouldn't happen. Our ModelToolCall.call_func actually passes in everything as kwargs.
+                        FancyLogger.get_logger().warning(
+                            f"ignoring unexpected args while calling langchain tool ({tool_name}): ({args})"
+                        )
+
+                    return tool.run(tool_input={**kwargs})
+
+                tool_call = parameter_remapper
+                return MelleaTool(tool_name, tool_call, as_json)
+            else:
+                raise ValueError(
+                    f"tool parameter must be a langchain tool type; got: {type(tool)}"
+                )
+
+        except ImportError as e:
+            raise ImportError(
+                f"It appears you are attempting to utilize a langchain tool '{type(tool)}'"
+                "Please install langchain core: pip install 'langchain_core'."
+            ) from e
+
+    @classmethod
+    def from_callable(cls, func: Callable, name: str | None = None):
+        """Create a Tool from a callable."""
+        # Use the function name if the name is '' or None.
+        tool_name = name or func.__name__
+        as_json = convert_function_to_ollama_tool(func, tool_name).model_dump(
+            exclude_none=True
+        )
+        tool_call = func
+        return MelleaTool(tool_name, tool_call, as_json)
+
+
 def add_tools_from_model_options(
-    tools_dict: dict[str, Callable], model_options: dict[str, Any]
+    tools_dict: dict[str, AbstractMelleaTool], 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)
@@ -23,30 +99,31 @@ def add_tools_from_model_options(
 
     # Mappings are iterable.
     assert isinstance(model_opts_tools, Iterable), (
-        "ModelOption.TOOLS must be a list of Callables or dict[str, Callable]"
+        "ModelOption.TOOLS must be a list of Tool or dict[str, Tool]"
     )
 
     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"
+        for tool_name, tool_instance in model_opts_tools.items():
+            assert isinstance(tool_name, str), (
+                f"If ModelOption.TOOLS is a dict, it must be a dict of [str, Tool]; found {type(tool_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"
+            assert isinstance(tool_instance, MelleaTool), (
+                f"If ModelOption.TOOLS is a dict, it must be a dict of [str, Tool]; found {type(tool_instance)} as the value instead"
             )
-            tools_dict[func_name] = func
+            tools_dict[tool_name] = tool_instance
     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)}"
+        for tool_instance in model_opts_tools:
+            assert isinstance(tool_instance, MelleaTool), (
+                f"If ModelOption.TOOLS is a list, it must be a list of Tool; found {type(tool_instance)}"
             )
-            tools_dict[func.__name__] = func
+            tools_dict[tool_instance.name] = tool_instance
 
 
 def add_tools_from_context_actions(
-    tools_dict: dict[str, Callable], ctx_actions: list[Component | CBlock] | None
+    tools_dict: dict[str, AbstractMelleaTool],
+    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:
@@ -64,24 +141,15 @@ def add_tools_from_context_actions(
             tools_dict[tool_name] = func
 
 
-def convert_tools_to_json(tools: dict[str, Callable]) -> list[dict]:
+def convert_tools_to_json(tools: dict[str, AbstractMelleaTool]) -> list[dict]:
     """Convert tools to json dict representation.
 
     Notes:
     - Huggingface transformers library lets you pass in an array of functions but doesn't like methods.
     - WatsonxAI uses `from langchain_ibm.chat_models import convert_to_openai_tool` in their demos, but it gives the same values.
     - OpenAI uses the same format / schema.
     """
-    converted: list[dict[str, Any]] = []
-    for tool in tools.values():
-        try:
-            converted.append(
-                convert_function_to_tool(tool).model_dump(exclude_none=True)
-            )
-        except Exception:
-            pass
-
-    return converted
+    return [t.as_json_tool for t in tools.values()]
 
 
 def json_extraction(text: str) -> Generator[dict, None, None]:
@@ -209,7 +277,7 @@ def __contains__(self, key: str) -> bool:
         >>> msg['tool_calls'] = None
         >>> 'tool_calls' in msg
         True
-        >>> tool = Tool()
+        >>> tool = OllamaTool()
         >>> 'type' in tool
         True
         """
@@ -240,7 +308,7 @@ def get(self, key: str, default: Any = None) -> Any:
 
 
 # https://github.com/ollama/ollama-python/blob/60e7b2f9ce710eeb57ef2986c46ea612ae7516af/ollama/_types.py#L337-L363
-class Tool(SubscriptableBaseModel):
+class OllamaTool(SubscriptableBaseModel):
     """Class imported from Ollama."""
 
     type: str | None = "function"
@@ -323,7 +391,9 @@ def _parse_docstring(doc_string: str | None) -> dict[str, str]:
 
 
 # https://github.com/ollama/ollama-python/blob/60e7b2f9ce710eeb57ef2986c46ea612ae7516af/ollama/_utils.py#L56-L90
-def convert_function_to_tool(func: Callable) -> Tool:
+def convert_function_to_ollama_tool(
+    func: Callable, name: str | None = None
+) -> OllamaTool:
     """Imported from Ollama."""
     doc_string_hash = str(hash(inspect.getdoc(func)))
     parsed_docstring = _parse_docstring(inspect.getdoc(func))
@@ -356,13 +426,13 @@ def convert_function_to_tool(func: Callable) -> Tool:
             "type": ", ".join(types),
         }
 
-    tool = Tool(
+    tool = OllamaTool(
         type="function",
-        function=Tool.Function(
-            name=func.__name__,
+        function=OllamaTool.Function(
+            name=name or func.__name__,
             description=schema.get("description", ""),
-            parameters=Tool.Function.Parameters(**schema),
+            parameters=OllamaTool.Function.Parameters(**schema),
         ),
     )
 
-    return Tool.model_validate(tool)
+    return OllamaTool.model_validate(tool)

mellea/backends/utils.py

@@ -6,6 +6,7 @@
 from collections.abc import Callable
 
 from ..core import CBlock, Component, Context, FancyLogger, ModelToolCall
+from ..core.base import AbstractMelleaTool
 from ..formatters import ChatFormatter
 from ..stdlib.components import Message
 from .tools import parse_tools
@@ -56,7 +57,7 @@ def to_chat(
 
 
 def to_tool_calls(
-    tools: dict[str, Callable], decoded_result: str
+    tools: dict[str, AbstractMelleaTool], decoded_result: str
 ) -> dict[str, ModelToolCall] | None:
     """Parse a tool call string."""
     model_tool_calls: dict[str, ModelToolCall] = dict()
@@ -70,8 +71,8 @@ def to_tool_calls(
 
         # Clean up the function args slightly. Some models seem to
         # hallucinate parameters when none are required.
-        sig = inspect.signature(func)
-        if len(sig.parameters) == 0:
+        param_map = func.as_json_tool["function"]["parameters"]["properties"]
+        if len(param_map) == 0:
             tool_args = {}
 
         model_tool_calls[tool_name] = ModelToolCall(tool_name, func, tool_args)