hub/backend/src/tool.py

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

    # -- 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 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=None, id=None):
        self.id = id or ""
        self.tool_name = tool_name
        self.arguments = arguments or {}

        self.result = None

    def __repr__(self):
        return f"ToolCall(tool={self.tool_name!r}, args={self.arguments!r})"
    
    def call_message(self):
        raise Exception("TODO: Llama message for tool call")
    
    def result_message(self):
        from inference import LlamaMessage
        return LlamaMessage("tool", self.result)


# ---------------------------------------------------------------------------
# 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)