266 lines
9.2 KiB
Python
266 lines
9.2 KiB
Python
import inspect
|
|
import asyncio
|
|
from typing import get_type_hints
|
|
|
|
class ToolParameter:
|
|
"""Describes a single input parameter for a tool.
|
|
|
|
Mirrors the JSON-schema style used by OpenAI / llama-cpp function calling.
|
|
|
|
Attributes:
|
|
name: Parameter name.
|
|
type: JSON-schema type string, e.g. "string", "number", "boolean".
|
|
description: Human-readable description of the parameter.
|
|
required: Whether the parameter is required.
|
|
enum: Optional list of allowed values.
|
|
"""
|
|
|
|
def __init__(self, name, type="string", description="", required=True, enum=None):
|
|
self.name = name
|
|
self.type = type
|
|
self.description = description
|
|
self.required = required
|
|
self.enum = enum
|
|
|
|
def to_dict(self):
|
|
d = {
|
|
"type": self.type,
|
|
"description": self.description,
|
|
}
|
|
if self.enum is not None:
|
|
d["enum"] = self.enum
|
|
return d
|
|
|
|
def __repr__(self):
|
|
return f"ToolParameter(name={self.name!r}, type={self.type!r})"
|
|
|
|
|
|
class ToolSpec:
|
|
"""Declarative specification for a tool callable by an LLM.
|
|
|
|
A ToolSpec describes *what* a tool does and *what* parameters it accepts.
|
|
It is used at inference time to build the tool-description payload sent to
|
|
the model so the model can emit structured tool-call requests.
|
|
|
|
Inspired by ``langchain.tools.StructuredTool`` and OpenAI's
|
|
``tools`` / ``functions`` request format.
|
|
|
|
Attributes:
|
|
name: Unique identifier for the tool.
|
|
description: Short human-readable description shown to the model.
|
|
parameters: List of ``ToolParameter`` objects.
|
|
# execution
|
|
func: Optional callable that executes the tool. When present the
|
|
orchestrator can dispatch a ``ToolCall`` automatically.
|
|
"""
|
|
|
|
def __init__(self, name, description="", parameters=None, func=None):
|
|
self.name = name
|
|
self.description = description
|
|
self.parameters = parameters or []
|
|
self.func = func
|
|
|
|
def clone (self):
|
|
return ToolSpec(
|
|
name = self.name,
|
|
description = self.description,
|
|
parameters = self.parameters,
|
|
func = self.func,
|
|
)
|
|
|
|
# -- serialization helpers -------------------------------------------
|
|
|
|
def to_openai_dict(self):
|
|
"""Return a dict in OpenAI chat-completions ``tools`` format."""
|
|
properties = {}
|
|
required = []
|
|
for p in self.parameters:
|
|
properties[p.name] = p.to_dict()
|
|
if p.required:
|
|
required.append(p.name)
|
|
|
|
return {
|
|
"type": "function",
|
|
"function": {
|
|
"name": self.name,
|
|
"description": self.description,
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": properties,
|
|
"required": required,
|
|
},
|
|
},
|
|
}
|
|
|
|
def to_replicate_dict(self):
|
|
"""Return a dict suitable for Replicate model input.
|
|
|
|
Replicate models do not have a standard tool-calling protocol, so we
|
|
return a simplified representation that can be injected into the
|
|
prompt as text.
|
|
"""
|
|
param_lines = []
|
|
for p in self.parameters:
|
|
req = " (required)" if p.required else " (optional)"
|
|
param_lines.append(f" - {p.name} ({p.type}){req}: {p.description}")
|
|
params_text = "\n".join(param_lines) if param_lines else " (none)"
|
|
return {
|
|
"name": self.name,
|
|
"description": self.description,
|
|
"parameters_text": f"{self.name}:\n{params_text}",
|
|
}
|
|
|
|
def execute(self, **kwargs):
|
|
"""Synchronously run the tool with the given keyword arguments."""
|
|
if self.func is None:
|
|
raise RuntimeError(f"ToolSpec '{self.name}' has no callable (func is None)")
|
|
return self.func(**kwargs)
|
|
|
|
async def execute_async(self, **kwargs):
|
|
"""Asynchronously run the tool with the given keyword arguments."""
|
|
if self.func is None:
|
|
raise RuntimeError(f"ToolSpec '{self.name}' has no callable (func is None)")
|
|
if asyncio.iscoroutinefunction(self.func):
|
|
return await self.func(**kwargs)
|
|
return self.func(**kwargs)
|
|
|
|
def __repr__(self):
|
|
return f"ToolSpec(name={self.name!r}, params={[p.name for p in self.parameters]})"
|
|
|
|
class Tool:
|
|
"""Base class for tools that implement their logic in a 'run' method.
|
|
|
|
Subclasses should override 'run' to provide the tool's functionality.
|
|
|
|
The Tool class provides a method to_spec() that automatically generates
|
|
a ToolSpec from the class name, docstring, and run method signature.
|
|
This allows tools to be used with LLM function calling without manually
|
|
creating ToolSpec objects. The generated ToolSpec is cached at initialization.
|
|
"""
|
|
|
|
def __init__(self, name=None):
|
|
# Generate and cache the ToolSpec during initialization
|
|
self.name = name or self.__class__.tool_name()
|
|
self._spec = self._generate_spec()
|
|
|
|
def run(self, **kwargs):
|
|
"""Execute the tool. To be overridden by subclasses."""
|
|
raise NotImplementedError("Tool subclasses must override the 'run' method.")
|
|
|
|
@classmethod
|
|
def tool_name(cls):
|
|
return ''.join(['-'+c.lower() if c.isupper() else c for c in cls.__name__]).lstrip('_')
|
|
|
|
|
|
def _generate_spec(self):
|
|
"""Generate a ToolSpec from this tool instance.
|
|
|
|
Automatically creates a ToolSpec by:
|
|
- Using the class name (converted to snake_case) as the tool name
|
|
- Using the class docstring as the description
|
|
- Inspecting the run method signature to create parameters
|
|
"""
|
|
|
|
# Use the class docstring as the description
|
|
description = inspect.getdoc(self.__class__) or ""
|
|
|
|
# Get parameters from the run method signature
|
|
parameters = []
|
|
run_signature = inspect.signature(self.run)
|
|
type_hints = get_type_hints(self.run)
|
|
|
|
for param_name, param in run_signature.parameters.items():
|
|
# Skip 'self' parameter
|
|
if param_name == 'self':
|
|
continue
|
|
|
|
# Get type information
|
|
param_type = type_hints.get(param_name, "string")
|
|
if isinstance(param_type, type):
|
|
type_name = param_type.__name__.lower()
|
|
else:
|
|
type_name = str(param_type).lower()
|
|
|
|
# Check if parameter is required (no default or default is None)
|
|
is_required = param.default is inspect.Parameter.empty or param.default is None
|
|
|
|
# Create ToolParameter
|
|
tool_param = ToolParameter(
|
|
name=param_name,
|
|
type=type_name,
|
|
description="",
|
|
required=is_required
|
|
)
|
|
parameters.append(tool_param)
|
|
|
|
return ToolSpec(
|
|
name=self.name,
|
|
description=description,
|
|
parameters=parameters,
|
|
func=self.run,
|
|
)
|
|
|
|
def get_spec(self):
|
|
"""Return the cached ToolSpec for this tool."""
|
|
return self._spec
|
|
|
|
def __repr__(self):
|
|
return f"Tool(name={self._spec.name!r}, params={[p.name for p in self._spec.parameters]})"
|
|
|
|
class ToolCall:
|
|
"""Represents a single tool invocation request emitted by the model.
|
|
|
|
Attributes:
|
|
id: Unique call identifier (used by OpenAI / llama-cpp protocol).
|
|
tool_name: Name of the tool to invoke.
|
|
arguments: Dict of argument name → value.
|
|
result: String content returned by the tool.
|
|
is_error: Whether the result represents an error.
|
|
"""
|
|
|
|
def __init__(self, tool_name, arguments={}, id=None, result=None):
|
|
self.id = id or ""
|
|
self.tool_name = tool_name
|
|
self.arguments = arguments
|
|
|
|
self.result = result
|
|
|
|
def __repr__(self):
|
|
id_prefix = self.id + ', ' if self.id else ''
|
|
return f"ToolCall({id_prefix}tool={self.tool_name!r}, args={self.arguments!r})"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers to format tool specs per engine
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def format_tools_for_engine(tools, engine):
|
|
"""Convert a list of ``ToolSpec`` objects for the given engine.
|
|
|
|
Returns a value suitable to pass as the ``tools`` kwarg to the underlying
|
|
chat-completion call.
|
|
"""
|
|
specs = [t for t in tools]
|
|
|
|
if engine in ("llama-cpp", "ollama"):
|
|
return [s.to_openai_dict() for s in specs]
|
|
elif engine == "replicate":
|
|
return [s.to_replicate_dict() for s in specs]
|
|
else:
|
|
return [s.to_openai_dict() for s in specs]
|
|
|
|
|
|
def build_tool_prompt_prefix(tools):
|
|
"""Build a text prefix describing available tools for engines that rely
|
|
on prompt-injection (e.g. Replicate) rather than native tool-calling."""
|
|
if not tools:
|
|
return ""
|
|
parts = ["You have access to the following tools:"]
|
|
for t in tools:
|
|
spec = t
|
|
info = spec.to_replicate_dict()
|
|
parts.append(info["parameters_text"])
|
|
parts.append("")
|
|
return "\n".join(parts)
|
|
|