Add ToolKit for new interface

mirascope/core/base/__init__.py

@@ -10,10 +10,10 @@
 from .streams import BaseAsyncStream, BaseStream
 from .structured_streams import BaseAsyncStructuredStream, BaseStructuredStream
 from .tools import BaseTool
+from .toolkit import BaseToolKit, toolkit_tool
 
 __all__ = [
     "BaseAsyncStream",
-    "BaseAsyncStructuredStream",
     "BaseCallParams",
     "BaseCallResponse",
     "BaseCallResponseChunk",
@@ -23,7 +23,9 @@
     "BaseStream",
     "BaseStructuredStream",
     "BaseTool",
-    "_partial",
+    "BaseToolKit",
     "tags",
+    "toolkit_tool",
+    "_partial",
     "_utils",
 ]

mirascope/core/base/_utils.py

@@ -5,7 +5,6 @@
 from abc import update_abstractmethods
 from enum import Enum
 from string import Formatter
-from textwrap import dedent
 from typing import (
     Annotated,
     Any,
@@ -32,17 +31,19 @@
 """
 
 
-def format_prompt_template(template: str, attrs: dict[str, Any]) -> str:
-    """Formats the given prompt `template`"""
-    dedented_template = dedent(template).strip()
-    template_vars = [
-        var for _, var, _, _ in Formatter().parse(dedented_template) if var is not None
-    ]
+def get_template_variables(template: str) -> list[str]:
+    """Returns the variables in the given template string."""
+    return [var for _, var, _, _ in Formatter().parse(template) if var is not None]
+
 
+def get_template_values(
+    template_variables: list[str], attrs: dict[str, Any]
+) -> dict[str, Any]:
+    """Returns the values of the given `template_variables` from the provided `attrs`."""
     values = {}
     if "self" in attrs:
         values["self"] = attrs.get("self")
-    for var in template_vars:
+    for var in template_variables:
         if var.startswith("self"):
             continue
         attr = attrs[var]
@@ -57,6 +58,15 @@ def format_prompt_template(template: str, attrs: dict[str, Any]) -> str:
                 values[var] = "\n".join([str(item) for item in attr])
         else:
             values[var] = str(attr)
+    return values
+
+
+def format_template(template: str, attrs: dict[str, Any]) -> str:
+    """Formats the given prompt `template`"""
+    dedented_template = inspect.cleandoc(template).strip()
+    template_vars = get_template_variables(dedented_template)
+
+    values = get_template_values(template_vars, attrs)
 
     return dedented_template.format(**values)
 
@@ -94,14 +104,14 @@ def parse_prompt_messages(
                 )
             messages += attr
         else:
-            content = format_prompt_template(match.group(2), attrs)
+            content = format_template(match.group(2), attrs)
             if content:
                 messages.append({"role": role, "content": content})
     if len(messages) == 0:
         messages.append(
             {
                 "role": "user",
-                "content": format_prompt_template(template, attrs),
+                "content": format_template(template, attrs),
             }
         )
     return messages
@@ -113,7 +123,7 @@ def parse_prompt_messages(
 
 
 def convert_function_to_base_tool(
-    fn: Callable, base: type[BaseToolT]
+    fn: Callable, base: type[BaseToolT], __doc__: str | None = None
 ) -> type[BaseToolT]:
     """Constructst a `BaseToolT` type from the given function.
 
@@ -124,6 +134,7 @@ def convert_function_to_base_tool(
     Args:
         fn: The function to convert.
         base: The `BaseToolT` type to which the function is converted.
+        __doc__: The docstring to use for the constructed `BaseToolT` type.
 
     Returns:
         The constructed `BaseToolT` type.
@@ -136,13 +147,16 @@ def convert_function_to_base_tool(
             doesn't have a docstring description.
     """
     docstring = None
-    if fn.__doc__:
-        docstring = parse(fn.__doc__)
+    func_doc = __doc__ or fn.__doc__
+    if func_doc:
+        docstring = parse(func_doc)
 
     field_definitions = {}
     hints = get_type_hints(fn)
+    has_self_or_cls = False
     for i, parameter in enumerate(inspect.signature(fn).parameters.values()):
         if parameter.name == "self" or parameter.name == "cls":
+            has_self_or_cls = True
             continue
         if parameter.annotation == inspect.Parameter.empty:
             raise ValueError("All parameters must have a type annotation.")
@@ -181,20 +195,23 @@ def convert_function_to_base_tool(
     model = create_model(
         fn.__name__,
         __base__=base,
-        __doc__=inspect.cleandoc(fn.__doc__) if fn.__doc__ else DEFAULT_TOOL_DOCSTRING,
+        __doc__=inspect.cleandoc(func_doc) if func_doc else DEFAULT_TOOL_DOCSTRING,
         **cast(dict[str, Any], field_definitions),
     )
 
     def call(self: base):
         return fn(
-            **{
-                str(
-                    self.model_fields[field_name].alias
-                    if self.model_fields[field_name].alias
-                    else field_name
-                ): getattr(self, field_name)
-                for field_name in self.model_dump(exclude={"tool_call"})
-            }
+            **(
+                ({"self": self} if has_self_or_cls else {})
+                | {
+                    str(
+                        self.model_fields[field_name].alias
+                        if self.model_fields[field_name].alias
+                        else field_name
+                    ): getattr(self, field_name)
+                    for field_name in self.model_dump(exclude={"tool_call"})
+                }
+            )
         )
 
     setattr(model, "call", call)

mirascope/core/base/prompts.py

@@ -41,7 +41,7 @@ class BookRecommendationPrompt(BasePrompt):
 
     def __str__(self) -> str:
         """Returns the formatted template."""
-        return _utils.format_prompt_template(self.prompt_template, self.model_dump())
+        return _utils.format_template(self.prompt_template, self.model_dump())
 
     def message_params(self) -> list[BaseMessageParam]:
         """Returns the template as a formatted list of `Message` instances."""

mirascope/core/base/toolkit.py

@@ -0,0 +1,81 @@
+from __future__ import annotations
+
+import inspect
+from abc import ABC
+from typing import Callable, ClassVar, ParamSpec, NamedTuple
+
+from pydantic import BaseModel, ConfigDict
+from typing_extensions import LiteralString
+
+from . import BaseTool
+from ._utils import convert_function_to_base_tool, get_template_variables
+
+_TOOLKIT_TOOL_METHOD_MARKER: LiteralString = "__toolkit_tool_method__"
+
+P = ParamSpec("P")
+
+
+def toolkit_tool(
+    method: Callable[[BaseToolKit, ...], str],
+) -> Callable[[BaseToolKit, ...], str]:
+    # Mark the method as a toolkit tool
+    setattr(method, _TOOLKIT_TOOL_METHOD_MARKER, True)
+
+    return method
+
+
+class TookKitToolMethod(NamedTuple):
+    method: Callable[[BaseToolKit, ...], str]
+    template_vars: list[str]
+    template: str
+
+
+class BaseToolKit(BaseModel, ABC):
+    """A class for defining tools for LLM call tools."""
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+    _toolkit_tool_methods: ClassVar[list[TookKitToolMethod]]
+
+    def create_tools(self) -> list[type[BaseTool]]:
+        """The method to create the tools."""
+        tools: list[type[BaseTool]] = []
+        for method, template_vars, template in self._toolkit_tool_methods:
+            formatted_template = template.format(self=self)
+            tool = convert_function_to_base_tool(method, BaseTool, formatted_template)
+            tools.append(tool)
+        return tools
+
+    @classmethod
+    def __pydantic_init_subclass__(cls, **kwargs):
+        cls._toolkit_tool_methods = []
+        for attr in cls.__dict__.values():
+            if not getattr(attr, _TOOLKIT_TOOL_METHOD_MARKER, False):
+                continue
+            # Validate the toolkit_tool_method
+            if (template := attr.__doc__) is None:
+                raise ValueError("The toolkit_tool method must have a docstring")
+
+            dedented_template = inspect.cleandoc(template)
+            template_vars = get_template_variables(dedented_template)
+
+            for var in template_vars:
+                if not var.startswith("self."):
+                    # Should be supported un-self variables?
+                    raise ValueError(
+                        "The toolkit_tool method must use self. prefix in template variables"
+                    )
+
+                self_var = var[5:]
+
+                # Expecting pydantic model fields or class attribute and property
+                if self_var in cls.model_fields or hasattr(cls, self_var):
+                    continue
+                raise ValueError(
+                    f"The toolkit_tool method template variable {var} is not found in the class"
+                )
+
+            cls._toolkit_tool_methods.append(
+                TookKitToolMethod(attr, template_vars, dedented_template)
+            )
+        if not cls._toolkit_tool_methods:
+            raise ValueError("No toolkit_tool method found")

tests/core/base/test_toolkit.py

@@ -0,0 +1,35 @@
+"""Tests for the `toolkit` module."""
+from typing import Literal
+
+from mirascope.core.base import BaseToolKit, toolkit_tool
+
+
+def test_toolkit() -> None:
+    """Tests the `BaseToolKit` class and the `toolkit_tool` decorator."""
+
+    class BookRecommendationToolKit(BaseToolKit):
+        """A toolkit for recommending books."""
+
+        reading_level: Literal["beginner", "advanced"]
+
+        @toolkit_tool
+        def format_book(self, title: str, author: str) -> str:
+            """Returns the title and author of a book nicely formatted.
+
+            Reading level: {self.reading_level}
+            """
+            return f"{title} by {author}"
+
+    toolkit = BookRecommendationToolKit(reading_level="beginner")
+    tools = toolkit.create_tools()
+    assert len(tools) == 1
+    tool = tools[0]
+    assert tool._name() == "format_book"
+    assert (
+        tool._description()
+        == "Returns the title and author of a book nicely formatted.\n\nReading level: beginner"
+    )
+    assert (
+        tool(title="The Name of the Wind", author="Rothfuss, Patrick").call()
+        == "The Name of the Wind by Rothfuss, Patrick"
+    )

tests/core/base/test_utils.py

@@ -6,7 +6,7 @@
 import pytest
 from pydantic import BaseModel
 
-from mirascope.core.base import BaseTool, _utils
+from mirascope.core.base import _utils, BaseTool
 
 
 def test_format_prompt_template() -> None:
@@ -32,7 +32,7 @@ def test_format_prompt_template() -> None:
         "genres": genres,
         "authors_and_books": authors_and_books,
     }
-    formatted_prompt_template = _utils.format_prompt_template(prompt_template, attrs)
+    formatted_prompt_template = _utils.format_template(prompt_template, attrs)
     assert (
         formatted_prompt_template
         == dedent(
@@ -114,8 +114,7 @@ def fn(model_name: str = "", self=None, cls=None) -> str:
 def test_convert_function_to_base_model_errors() -> None:
     """Tests the various `ValueErro` cases in `convert_function_to_base_model`."""
 
-    def empty(param) -> str:
-        ...  # pragma: no cover
+    def empty(param) -> str: ...  # pragma: no cover
 
     with pytest.raises(ValueError):
         _utils.convert_function_to_base_tool(empty, BaseTool)