Add ToolKit for new interface

mirascope/core/_internal/utils.py

@@ -28,15 +28,17 @@
 """
 
 
-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 = {}
-    for var in template_vars:
+    for var in template_variables:
         attr = attrs[var]
         if isinstance(attr, list):
             if len(attr) == 0:
@@ -49,6 +51,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_prompt_template(template: str, attrs: dict[str, Any]) -> str:
+    """Formats the given prompt `template`"""
+    dedented_template = dedent(template).strip()
+    template_vars = get_template_variables(dedented_template)
+
+    values = get_template_values(template_vars, attrs)
 
     return dedented_template.format(**values)
 

mirascope/core/base/__init__.py

@@ -3,5 +3,14 @@
 from .prompt import BasePrompt, tags
 from .tool import BaseTool
 from .types import BaseCallResponse, MessageParam
+from .toolkit import BaseToolKit, toolkit_tool
 
-__all__ = ["BaseCallResponse", "BasePrompt", "BaseTool", "MessageParam", "tags"]
+__all__ = [
+    "BaseCallResponse",
+    "BasePrompt",
+    "BaseTool",
+    "MessageParam",
+    "tags",
+    "BaseToolKit",
+    "toolkit_tool",
+]

mirascope/core/base/toolkit.py

@@ -0,0 +1,87 @@
+from __future__ import annotations
+
+from abc import ABC
+from textwrap import dedent
+from typing import Callable, Optional, ClassVar, ParamSpec
+from functools import wraps
+
+from pydantic import BaseModel, ConfigDict
+from typing_extensions import LiteralString
+
+from .tool import BaseTool
+from .._internal.utils import 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)
+
+    # TODO: Validate first argument is self
+    @wraps(method)
+    def inner(*args, **kwargs):
+        return method(*args, **kwargs)
+
+    return inner
+
+
+class BaseToolKit(BaseModel, ABC):
+    """A class for defining tools for LLM call tools."""
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+    _toolkit_tool_method: ClassVar[Callable[..., str]]
+    _toolkit_template_vars: ClassVar[list[str]]
+
+    def create_tool(self) -> type[BaseTool]:
+        """The method to create the tools."""
+        formated_template = self._toolkit_tool_method.__doc__.format(
+            **{var: getattr(self, var) for var in self._toolkit_template_vars}
+        )
+        # TODO: Generate the tool class with the formatted template
+        ...
+
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        toolkit_tool_method: Optional[Callable] = None
+        for name, value in cls.__dict__.items():
+            if getattr(value, _TOOLKIT_TOOL_METHOD_MARKER, False):
+                if toolkit_tool_method:
+                    raise ValueError("Only one toolkit_tool method is allowed")
+                toolkit_tool_method = value
+        if not toolkit_tool_method:
+            raise ValueError("No toolkit_tool method found")
+
+        # Validate the toolkit_tool_method
+        if (template := toolkit_tool_method.__doc__) is None:
+            raise ValueError("The toolkit_tool method must have a docstring")
+
+        dedented_template = dedent(template).strip()
+        if not (template_vars := get_template_variables(dedented_template)):
+            raise ValueError("The toolkit_tool method must have template variables")
+
+        if dedented_template != template:
+            toolkit_tool_method.__doc__ = 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
+            # TODO: Check attribute type such like callable, property, etc.
+            if self_var in cls.model_fields_set 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_method = toolkit_tool_method
+        cls._toolkit_template_vars = template_vars

mirascope/core/openai/openai_call.py

@@ -42,6 +42,17 @@ def decorator(fn: Callable[P, R]) -> Callable[P, OpenAICallResponse]:
         def call(*args: P.args, **kwargs: P.kwargs) -> OpenAICallResponse:
             prompt_template = inspect.getdoc(fn)
             assert prompt_template is not None, "The function must have a docstring."
+
+            # Try to get the dictionary for tools from the function result
+            # tools = []
+            fn_result = fn(*args, **kwargs)
+            if isinstance(fn_result, dict):
+                if fn_result_tools := fn_result.get("tools"):
+                    for fn_result_tool in fn_result_tools:
+                        # TODO: Generate tools
+                        # tools.append(generate(fn_result_tool))
+                        pass
+
             attrs = inspect.signature(fn).bind(*args, **kwargs).arguments
             messages = utils.parse_prompt_messages(
                 roles=["system", "user", "assistant"],