mars-mcp/mars_toolkit/core/llm_tools.py

"""
LLM Tools Module

This module provides decorators and utilities for defining, registering, and managing LLM tools.
It allows marking functions as LLM tools, generating JSON schemas for them, and retrieving
registered tools for use with LLM APIs.

"""

import asyncio
import inspect
import json
from functools import wraps
from typing import Any, Callable, Dict, List, Optional, get_type_hints, get_origin, get_args
import docstring_parser
from pydantic import BaseModel, create_model, Field

# Registry to store all registered tools
_TOOL_REGISTRY = {}

def llm_tool(name: Optional[str] = None, description: Optional[str] = None):
    """
    Decorator to mark a function as an LLM tool.

    This decorator registers the function as an LLM tool, generates a JSON schema for it,
    and makes it available for retrieval through the get_tools function.

    Args:
        name: Optional custom name for the tool. If not provided, the function name will be used.
        description: Optional custom description for the tool. If not provided, the function's
                    docstring will be used.

    Returns:
        The decorated function with additional attributes for LLM tool functionality.

    Example:
        @llm_tool(name="weather_lookup", description="Get current weather for a location")
        def get_weather(location: str, units: str = "metric") -> Dict[str, Any]:
            '''Get weather information for a specific location.'''
            # Implementation...
            return {"temperature": 22.5, "conditions": "sunny"}
    """
    # Handle case when decorator is used without parentheses: @llm_tool
    if callable(name):
        func = name
        name = None
        description = None
        return _llm_tool_impl(func, name, description)

    # Handle case when decorator is used with parentheses: @llm_tool() or @llm_tool(name="xyz")
    def decorator(func: Callable) -> Callable:
        return _llm_tool_impl(func, name, description)

    return decorator

def _llm_tool_impl(func: Callable, name: Optional[str] = None, description: Optional[str] = None) -> Callable:
    """Implementation of the llm_tool decorator."""
    # Get function signature and docstring
    sig = inspect.signature(func)
    doc = inspect.getdoc(func) or ""
    parsed_doc = docstring_parser.parse(doc)

    # Determine tool name
    tool_name = name or func.__name__

    # Determine tool description
    tool_description = description or doc

    # Create parameter properties for JSON schema
    properties = {}
    required = []

    for param_name, param in sig.parameters.items():
        # Skip self parameter for methods
        if param_name == "self":
            continue

        param_type = param.annotation
        param_default = None if param.default is inspect.Parameter.empty else param.default
        param_required = param.default is inspect.Parameter.empty

        # Get parameter description from docstring if available
        param_desc = ""
        for param_doc in parsed_doc.params:
            if param_doc.arg_name == param_name:
                param_desc = param_doc.description
                break

        # Handle Annotated types
        if get_origin(param_type) is not None and get_origin(param_type).__name__ == "Annotated":
            args = get_args(param_type)
            param_type = args[0]  # The actual type
            if len(args) > 1 and isinstance(args[1], str):
                param_desc = args[1]  # The description

        # Create property for parameter
        param_schema = {
            "type": _get_json_type(param_type),
            "description": param_desc,
            "title": param_name.replace("_", " ").title()
        }

        # Add default value if available
        if param_default is not None:
            param_schema["default"] = param_default

        properties[param_name] = param_schema

        # Add to required list if no default value
        if param_required:
            required.append(param_name)

    # Create JSON schema
    schema = {
        "type": "function",
        "function": {
            "name": tool_name,
            "description": tool_description,
            "parameters": {
                "type": "object",
                "properties": properties,
                "required": required
            }
        }
    }

    # Create Pydantic model for args schema
    field_definitions = {}
    for param_name, param in sig.parameters.items():
        if param_name == "self":
            continue

        param_type = param.annotation
        param_default = ... if param.default is inspect.Parameter.empty else param.default

        # Handle Annotated types
        if get_origin(param_type) is not None and get_origin(param_type).__name__ == "Annotated":
            args = get_args(param_type)
            param_type = args[0]
            description = args[1] if len(args) > 1 and isinstance(args[1], str) else ""
            field_definitions[param_name] = (param_type, Field(default=param_default, description=description))
        else:
            field_definitions[param_name] = (param_type, Field(default=param_default))

    # Create args schema model
    model_name = f"{tool_name.title().replace('_', '')}Schema"
    args_schema = create_model(model_name, **field_definitions)

    # 根据原始函数是否是异步函数来创建相应类型的包装函数
    if asyncio.iscoroutinefunction(func):
        @wraps(func)
        async def wrapper(*args, **kwargs):
            return await func(*args, **kwargs)
    else:
        @wraps(func)
        def wrapper(*args, **kwargs):
            return func(*args, **kwargs)

    # Attach metadata to function
    wrapper.is_llm_tool = True
    wrapper.tool_name = tool_name
    wrapper.tool_description = tool_description
    wrapper.json_schema = schema
    wrapper.args_schema = args_schema

    # Register the tool
    _TOOL_REGISTRY[tool_name] = wrapper

    return wrapper

def get_tools() -> Dict[str, Callable]:
    """
    Get all registered LLM tools.

    Returns:
        A dictionary mapping tool names to their corresponding functions.
    """
    return _TOOL_REGISTRY

def get_tool_schemas() -> List[Dict[str, Any]]:
    """
    Get JSON schemas for all registered LLM tools.

    Returns:
        A list of JSON schemas for all registered tools, suitable for use with LLM APIs.
    """
    return [tool.json_schema for tool in _TOOL_REGISTRY.values()]

def _get_json_type(python_type: Any) -> str:
    """
    Convert Python type to JSON schema type.

    Args:
        python_type: Python type annotation

    Returns:
        Corresponding JSON schema type as string
    """
    if python_type is str:
        return "string"
    elif python_type is int:
        return "integer"
    elif python_type is float:
        return "number"
    elif python_type is bool:
        return "boolean"
    elif python_type is list or python_type is List:
        return "array"
    elif python_type is dict or python_type is Dict:
        return "object"
    else:
        # Default to string for complex types
        return "string"