tools.py

# Copyright (c) Microsoft Corporation. All rights reserved.

"""
Tool definition utilities for the Copilot SDK.

Provides a decorator-based API for defining tools with automatic JSON schema
generation from Pydantic models.
"""

from __future__ import annotations

import inspect
import json
from collections.abc import Awaitable, Callable
from dataclasses import dataclass, field
from datetime import date, datetime, time
from decimal import Decimal
from enum import Enum
from typing import TYPE_CHECKING, Any, Literal, TypeVar, get_type_hints, overload
from uuid import UUID

from pydantic import BaseModel, ValidationError

if TYPE_CHECKING:
    from .generated.rpc import CurrentToolMetadata

from .generated.rpc import (
    ExternalToolTextResultForLlm,
    ExternalToolTextResultForLlmBinaryResultsForLlm,
    ExternalToolTextResultForLlmBinaryResultsForLlmType,
)

ToolResultType = Literal["success", "failure", "rejected", "denied", "timeout"]


@dataclass
class ToolBinaryResult:
    """Binary content returned by a tool."""

    data: str = ""
    mime_type: str = ""
    type: Literal["image", "resource"] = "image"
    description: str = ""


@dataclass
class ToolResult:
    """Result of a tool invocation."""

    text_result_for_llm: str = ""
    result_type: ToolResultType = "success"
    error: str | None = None
    binary_results_for_llm: list[ToolBinaryResult] | None = None
    session_log: str | None = None
    tool_telemetry: dict[str, Any] | None = None
    tool_references: list[str] | None = None
    _from_exception: bool = field(default=False, repr=False)


@dataclass
class ToolInvocation:
    """Context passed to a tool handler when invoked."""

    session_id: str = ""
    tool_call_id: str = ""
    tool_name: str = ""
    arguments: Any = None
    available_tools: list[CurrentToolMetadata] | None = None
    """Snapshot of the session's currently initialized tools.

    Populated by the SDK only when this invocation targets the built-in
    tool-search tool (``tool_search_tool``), so a tool-search override can
    rank/filter the live catalog -- including MCP tools configured in settings --
    without issuing its own RPC. ``None`` for every other tool invocation.
    """


ToolHandler = Callable[[ToolInvocation], ToolResult | Awaitable[ToolResult]]


@dataclass
class Tool:
    name: str
    description: str
    handler: ToolHandler | None = None
    parameters: dict[str, Any] | None = None
    overrides_built_in_tool: bool = False
    skip_permission: bool = False
    defer: Literal["auto", "never"] | None = None
    metadata: dict[str, Any] | None = None
    #: When true, a successful call to this tool ends the agent turn: the
    #: runtime halts instead of feeding the result back to the model for
    #: another round. A failed call leaves the loop running so the model can
    #: read the error and retry.
    is_terminal: bool = False


T = TypeVar("T", bound=BaseModel)
R = TypeVar("R")


@overload
def define_tool(
    name: str | None = None,
    *,
    description: str | None = None,
    overrides_built_in_tool: bool = False,
    skip_permission: bool = False,
    defer: Literal["auto", "never"] | None = None,
    metadata: dict[str, Any] | None = None,
    is_terminal: bool = False,
) -> Callable[[Callable[..., Any]], Tool]:
    pass


@overload
def define_tool(
    name: str,
    *,
    description: str | None = None,
    params_type: type[T],
    handler: None = None,
    overrides_built_in_tool: bool = False,
    skip_permission: bool = False,
    defer: Literal["auto", "never"] | None = None,
    metadata: dict[str, Any] | None = None,
    is_terminal: bool = False,
) -> Tool:
    pass


@overload
def define_tool(
    name: str,
    *,
    description: str | None = None,
    handler: Callable[[T, ToolInvocation], R],
    params_type: type[T],
    overrides_built_in_tool: bool = False,
    skip_permission: bool = False,
    defer: Literal["auto", "never"] | None = None,
    metadata: dict[str, Any] | None = None,
    is_terminal: bool = False,
) -> Tool:
    pass


def define_tool(
    name: str | None = None,
    *,
    description: str | None = None,
    handler: Callable[[Any, ToolInvocation], Any] | None = None,
    params_type: type[BaseModel] | None = None,
    overrides_built_in_tool: bool = False,
    skip_permission: bool = False,
    defer: Literal["auto", "never"] | None = None,
    metadata: dict[str, Any] | None = None,
    is_terminal: bool = False,
) -> Tool | Callable[[Callable[[Any, ToolInvocation], Any]], Tool]:
    """
    Define a tool with automatic JSON schema generation from Pydantic models.

    Can be used as a decorator or as a function:

    Decorator usage (recommended):

        from pydantic import BaseModel, Field

        class LookupIssueParams(BaseModel):
            id: str = Field(description="Issue identifier")

        @define_tool(description="Fetch issue details")
        def lookup_issue(params: LookupIssueParams) -> str:
            return fetch_issue(params.id).summary

    Function usage:

        tool = define_tool(
            "lookup_issue",
            description="Fetch issue details",
            handler=lambda params, inv: fetch_issue(params.id).summary,
            params_type=LookupIssueParams
        )

    Declaration-only usage:

        tool = define_tool(
            "lookup_issue",
            description="Fetch issue details",
            params_type=LookupIssueParams,
        )

    Args:
        name: The tool name (defaults to function name)
        description: Description of what the tool does (shown to the LLM)
        handler: Optional handler function (if not using as decorator)
        params_type: Optional Pydantic model type for parameters (inferred from
                    type hints when using as decorator)
        overrides_built_in_tool: When True, explicitly indicates this tool is intended
                    to override a built-in tool of the same name. If not set and the
                    name clashes with a built-in tool, the runtime will return an error.
        skip_permission: When True, the tool can execute without a permission prompt.
        defer: Controls whether the tool may be deferred (loaded lazily via tool search)
                    rather than always pre-loaded. When "auto", the tool can be deferred
                    and surfaced through tool search. When "never", the tool is always
                    pre-loaded. Optional; defaults to "auto".
        metadata: Opaque, host-defined metadata associated with the tool definition.
                    Keys are namespaced and not part of the stable public API; values
                    are not interpreted and may be recognized to inform host-specific
                    behavior. Unknown keys are preserved.
        is_terminal: When True, a successful call to this tool ends the agent turn:
                    the runtime halts instead of feeding the result back to the model
                    for another round. A failed call leaves the loop running so the
                    model can read the error and retry.

    Returns:
        A Tool instance
    """

    def decorator(fn: Callable[..., Any]) -> Tool:
        tool_name = name if name is not None else getattr(fn, "__name__", "unknown")

        sig = inspect.signature(fn)
        param_names = list(sig.parameters.keys())
        hints = get_type_hints(fn)
        num_params = len(param_names)

        # Detect handler signature:
        # - 0 params: handler()
        # - 1 param, ToolInvocation: handler(invocation)
        # - 1 param, Pydantic: handler(params)
        # - 2 params: handler(params, invocation)
        ptype = params_type
        first_param_type = hints.get(param_names[0]) if param_names else None

        if num_params == 0:
            takes_params = False
            takes_invocation = False
        elif num_params == 1 and first_param_type is ToolInvocation:
            takes_params = False
            takes_invocation = True
        else:
            takes_params = True
            takes_invocation = num_params >= 2
            if ptype is None and _is_pydantic_model(first_param_type):
                ptype = first_param_type

        # Generate schema from Pydantic model
        schema = None
        if ptype is not None and _is_pydantic_model(ptype):
            schema = ptype.model_json_schema()

        async def wrapped_handler(invocation: ToolInvocation) -> ToolResult:
            try:
                # Build args based on detected signature
                call_args = []
                if takes_params:
                    args = invocation.arguments if invocation.arguments is not None else {}
                    if ptype is not None and _is_pydantic_model(ptype):
                        try:
                            call_args.append(ptype.model_validate(args))
                        except ValidationError as exc:
                            # Highlight input validation problems to the LLM.
                            parts = []
                            for err in exc.errors():
                                loc = ".".join(map(str, err["loc"]))
                                msg = err["msg"]
                                parts.append(f"{loc}: {msg}" if loc else msg)
                            return ToolResult(
                                text_result_for_llm="Invalid tool arguments:\n" + "\n".join(parts),
                                result_type="failure",
                                error=str(exc),
                                tool_telemetry={},
                            )
                    else:
                        call_args.append(args)
                if takes_invocation:
                    call_args.append(invocation)

                result = fn(*call_args)

                if inspect.isawaitable(result):
                    result = await result

                return _normalize_result(result)

            except Exception as exc:
                # Don't expose detailed error information to the LLM for security reasons.
                # The actual error is stored in the 'error' field for debugging.
                return ToolResult(
                    text_result_for_llm=(
                        "Invoking this tool produced an error. "
                        "Detailed information is not available."
                    ),
                    result_type="failure",
                    error=str(exc),
                    tool_telemetry={},
                    _from_exception=True,
                )

        return Tool(
            name=tool_name,
            description=description or "",
            parameters=schema,
            handler=wrapped_handler,
            overrides_built_in_tool=overrides_built_in_tool,
            skip_permission=skip_permission,
            defer=defer,
            metadata=metadata,
            is_terminal=is_terminal,
        )

    # If handler is provided, call decorator immediately
    if handler is not None:
        if name is None:
            raise ValueError("name is required when using define_tool with handler=")
        return decorator(handler)

    # If a parameter model is provided without a handler, expose a declaration-only tool.
    if name is not None and params_type is not None:
        schema = params_type.model_json_schema() if _is_pydantic_model(params_type) else None
        return Tool(
            name=name,
            description=description or "",
            parameters=schema,
            handler=None,
            overrides_built_in_tool=overrides_built_in_tool,
            skip_permission=skip_permission,
            defer=defer,
            metadata=metadata,
            is_terminal=is_terminal,
        )

    # Otherwise return decorator for @define_tool(...) usage
    return decorator


def _is_pydantic_model(t: Any) -> bool:
    """Check if a type is a Pydantic BaseModel subclass."""
    try:
        return isinstance(t, type) and issubclass(t, BaseModel)
    except TypeError:
        return False


def _normalize_result(result: Any) -> ToolResult:
    """
    Convert any return value to a ToolResult.

    - None returns empty success
    - Strings pass through directly
    - ToolResult passes through
    - Everything else gets JSON-serialized (with Pydantic support)
    """
    if result is None:
        return ToolResult(
            text_result_for_llm="",
            result_type="success",
        )

    # ToolResult dataclass passes through directly
    if isinstance(result, ToolResult):
        return result

    # Strings pass through directly
    if isinstance(result, str):
        return ToolResult(
            text_result_for_llm=result,
            result_type="success",
        )

    # Everything else gets JSON-serialized (with common Python and Pydantic values)
    def default(obj: Any) -> Any:
        if isinstance(obj, BaseModel):
            return obj.model_dump(mode="json")
        if isinstance(obj, (date, datetime, time)):
            return obj.isoformat()
        if isinstance(obj, (Decimal, UUID)):
            return str(obj)
        if isinstance(obj, Enum):
            return obj.value
        if isinstance(obj, set):
            return list(obj)
        raise TypeError(f"Object of type {type(obj).__name__} is not JSON serializable")

    try:
        json_str = json.dumps(result, default=default)
    except (TypeError, ValueError) as exc:
        raise TypeError(f"Failed to serialize tool result: {exc}") from exc

    return ToolResult(
        text_result_for_llm=json_str,
        result_type="success",
    )


def convert_mcp_call_tool_result(call_result: dict[str, Any]) -> ToolResult:
    """Convert an MCP CallToolResult dict into a ToolResult."""
    text_parts: list[str] = []
    binary_results: list[ToolBinaryResult] = []

    for block in call_result["content"]:
        block_type = block.get("type")
        if block_type == "text":
            text = block.get("text", "")
            if isinstance(text, str):
                text_parts.append(text)
        elif block_type == "image":
            data = block.get("data", "")
            mime_type = block.get("mimeType", "")
            if isinstance(data, str) and data and isinstance(mime_type, str):
                binary_results.append(
                    ToolBinaryResult(
                        data=data,
                        mime_type=mime_type,
                        type="image",
                    )
                )
        elif block_type == "resource":
            resource = block.get("resource", {})
            if not isinstance(resource, dict):
                continue
            text = resource.get("text")
            if isinstance(text, str) and text:
                text_parts.append(text)
            blob = resource.get("blob")
            if isinstance(blob, str) and blob:
                mime_type = resource.get("mimeType")
                if not isinstance(mime_type, str) or not mime_type:
                    mime_type = "application/octet-stream"
                uri = resource.get("uri", "")
                binary_results.append(
                    ToolBinaryResult(
                        data=blob,
                        mime_type=mime_type,
                        type="resource",
                        description=uri if isinstance(uri, str) else "",
                    )
                )

    return ToolResult(
        text_result_for_llm="\n".join(text_parts),
        result_type="failure" if call_result.get("isError") is True else "success",
        binary_results_for_llm=binary_results if binary_results else None,
    )


def tool_result_to_external_tool_text_result_for_llm(
    tool_result: ToolResult,
) -> ExternalToolTextResultForLlm:
    """Convert a ToolResult into the RPC payload sent to HandlePendingToolCall."""
    binary_results_for_llm = None
    if tool_result.binary_results_for_llm:
        binary_results_for_llm = [
            ExternalToolTextResultForLlmBinaryResultsForLlm(
                data=binary_result.data,
                mime_type=binary_result.mime_type,
                type=ExternalToolTextResultForLlmBinaryResultsForLlmType(binary_result.type),
                description=binary_result.description or None,
            )
            for binary_result in tool_result.binary_results_for_llm
        ]

    return ExternalToolTextResultForLlm(
        text_result_for_llm=tool_result.text_result_for_llm,
        binary_results_for_llm=binary_results_for_llm,
        error=tool_result.error,
        result_type=tool_result.result_type,
        session_log=tool_result.session_log,
        tool_references=tool_result.tool_references,
        tool_telemetry=tool_result.tool_telemetry,
    )