aphilippov/python-server-api

1import functools2import inspect3import json4from logging import getLogger5from typing import Any, Callable, Dict, ForwardRef, List, Optional, Set, Tuple, Type, TypeVar, Union6 7from pydantic import BaseModel, Field8from typing_extensions import Annotated, Literal, get_args, get_origin9 10from ._pydantic import JsonSchemaValue, evaluate_forwardref, model_dump, model_dump_json, type2schema11 12logger = getLogger(__name__)13 14T = TypeVar("T")15 16 17def get_typed_annotation(annotation: Any, globalns: Dict[str, Any]) -> Any:18    """Get the type annotation of a parameter.19 20    Args:21        annotation: The annotation of the parameter22        globalns: The global namespace of the function23 24    Returns:25        The type annotation of the parameter26    """27    if isinstance(annotation, str):28        annotation = ForwardRef(annotation)29        annotation = evaluate_forwardref(annotation, globalns, globalns)30    return annotation31 32 33def get_typed_signature(call: Callable[..., Any]) -> inspect.Signature:34    """Get the signature of a function with type annotations.35 36    Args:37        call: The function to get the signature for38 39    Returns:40        The signature of the function with type annotations41    """42    signature = inspect.signature(call)43    globalns = getattr(call, "__globals__", {})44    typed_params = [45        inspect.Parameter(46            name=param.name,47            kind=param.kind,48            default=param.default,49            annotation=get_typed_annotation(param.annotation, globalns),50        )51        for param in signature.parameters.values()52    ]53    typed_signature = inspect.Signature(typed_params)54    return typed_signature55 56 57def get_typed_return_annotation(call: Callable[..., Any]) -> Any:58    """Get the return annotation of a function.59 60    Args:61        call: The function to get the return annotation for62 63    Returns:64        The return annotation of the function65    """66    signature = inspect.signature(call)67    annotation = signature.return_annotation68 69    if annotation is inspect.Signature.empty:70        return None71 72    globalns = getattr(call, "__globals__", {})73    return get_typed_annotation(annotation, globalns)74 75 76def get_param_annotations(typed_signature: inspect.Signature) -> Dict[int, Union[Annotated[Type[Any], str], Type[Any]]]:77    """Get the type annotations of the parameters of a function78 79    Args:80        typed_signature: The signature of the function with type annotations81 82    Returns:83        A dictionary of the type annotations of the parameters of the function84    """85    return {86        k: v.annotation for k, v in typed_signature.parameters.items() if v.annotation is not inspect.Signature.empty87    }88 89 90class Parameters(BaseModel):91    """Parameters of a function as defined by the OpenAI API"""92 93    type: Literal["object"] = "object"94    properties: Dict[str, JsonSchemaValue]95    required: List[str]96 97 98class Function(BaseModel):99    """A function as defined by the OpenAI API"""100 101    description: Annotated[str, Field(description="Description of the function")]102    name: Annotated[str, Field(description="Name of the function")]103    parameters: Annotated[Parameters, Field(description="Parameters of the function")]104 105 106class ToolFunction(BaseModel):107    """A function under tool as defined by the OpenAI API."""108 109    type: Literal["function"] = "function"110    function: Annotated[Function, Field(description="Function under tool")]111 112 113def get_parameter_json_schema(114    k: str, v: Union[Annotated[Type[Any], str], Type[Any]], default_values: Dict[str, Any]115) -> JsonSchemaValue:116    """Get a JSON schema for a parameter as defined by the OpenAI API117 118    Args:119        k: The name of the parameter120        v: The type of the parameter121        default_values: The default values of the parameters of the function122 123    Returns:124        A Pydanitc model for the parameter125    """126 127    def type2description(k: str, v: Union[Annotated[Type[Any], str], Type[Any]]) -> str:128        # handles Annotated129        if hasattr(v, "__metadata__"):130            retval = v.__metadata__[0]131            if isinstance(retval, str):132                return retval133            else:134                raise ValueError(f"Invalid description {retval} for parameter {k}, should be a string.")135        else:136            return k137 138    schema = type2schema(v)139    if k in default_values:140        dv = default_values[k]141        schema["default"] = dv142 143    schema["description"] = type2description(k, v)144 145    return schema146 147 148def get_required_params(typed_signature: inspect.Signature) -> List[str]:149    """Get the required parameters of a function150 151    Args:152        signature: The signature of the function as returned by inspect.signature153 154    Returns:155        A list of the required parameters of the function156    """157    return [k for k, v in typed_signature.parameters.items() if v.default == inspect.Signature.empty]158 159 160def get_default_values(typed_signature: inspect.Signature) -> Dict[str, Any]:161    """Get default values of parameters of a function162 163    Args:164        signature: The signature of the function as returned by inspect.signature165 166    Returns:167        A dictionary of the default values of the parameters of the function168    """169    return {k: v.default for k, v in typed_signature.parameters.items() if v.default != inspect.Signature.empty}170 171 172def get_parameters(173    required: List[str],174    param_annotations: Dict[str, Union[Annotated[Type[Any], str], Type[Any]]],175    default_values: Dict[str, Any],176) -> Parameters:177    """Get the parameters of a function as defined by the OpenAI API178 179    Args:180        required: The required parameters of the function181        hints: The type hints of the function as returned by typing.get_type_hints182 183    Returns:184        A Pydantic model for the parameters of the function185    """186    return Parameters(187        properties={188            k: get_parameter_json_schema(k, v, default_values)189            for k, v in param_annotations.items()190            if v is not inspect.Signature.empty191        },192        required=required,193    )194 195 196def get_missing_annotations(typed_signature: inspect.Signature, required: List[str]) -> Tuple[Set[str], Set[str]]:197    """Get the missing annotations of a function198 199    Ignores the parameters with default values as they are not required to be annotated, but logs a warning.200    Args:201        typed_signature: The signature of the function with type annotations202        required: The required parameters of the function203 204    Returns:205        A set of the missing annotations of the function206    """207    all_missing = {k for k, v in typed_signature.parameters.items() if v.annotation is inspect.Signature.empty}208    missing = all_missing.intersection(set(required))209    unannotated_with_default = all_missing.difference(missing)210    return missing, unannotated_with_default211 212 213def get_function_schema(f: Callable[..., Any], *, name: Optional[str] = None, description: str) -> Dict[str, Any]:214    """Get a JSON schema for a function as defined by the OpenAI API215 216    Args:217        f: The function to get the JSON schema for218        name: The name of the function219        description: The description of the function220 221    Returns:222        A JSON schema for the function223 224    Raises:225        TypeError: If the function is not annotated226 227    Examples:228        ```229        def f(a: Annotated[str, "Parameter a"], b: int = 2, c: Annotated[float, "Parameter c"] = 0.1) -> None:230            pass231 232        get_function_schema(f, description="function f")233 234        #   {'type': 'function',235        #    'function': {'description': 'function f',236        #        'name': 'f',237        #        'parameters': {'type': 'object',238        #           'properties': {'a': {'type': 'str', 'description': 'Parameter a'},239        #               'b': {'type': 'int', 'description': 'b'},240        #               'c': {'type': 'float', 'description': 'Parameter c'}},241        #           'required': ['a']}}}242            ```243 244    """245    typed_signature = get_typed_signature(f)246    required = get_required_params(typed_signature)247    default_values = get_default_values(typed_signature)248    param_annotations = get_param_annotations(typed_signature)249    return_annotation = get_typed_return_annotation(f)250    missing, unannotated_with_default = get_missing_annotations(typed_signature, required)251 252    if return_annotation is None:253        logger.warning(254            f"The return type of the function '{f.__name__}' is not annotated. Although annotating it is "255            + "optional, the function should return either a string, a subclass of 'pydantic.BaseModel'."256        )257 258    if unannotated_with_default != set():259        unannotated_with_default_s = [f"'{k}'" for k in sorted(unannotated_with_default)]260        logger.warning(261            f"The following parameters of the function '{f.__name__}' with default values are not annotated: "262            + f"{', '.join(unannotated_with_default_s)}."263        )264 265    if missing != set():266        missing_s = [f"'{k}'" for k in sorted(missing)]267        raise TypeError(268            f"All parameters of the function '{f.__name__}' without default values must be annotated. "269            + f"The annotations are missing for the following parameters: {', '.join(missing_s)}"270        )271 272    fname = name if name else f.__name__273 274    parameters = get_parameters(required, param_annotations, default_values=default_values)275 276    function = ToolFunction(277        function=Function(278            description=description,279            name=fname,280            parameters=parameters,281        )282    )283 284    return model_dump(function)285 286 287def get_load_param_if_needed_function(t: Any) -> Optional[Callable[[T, Type[Any]], BaseModel]]:288    """Get a function to load a parameter if it is a Pydantic model289 290    Args:291        t: The type annotation of the parameter292 293    Returns:294        A function to load the parameter if it is a Pydantic model, otherwise None295 296    """297    if get_origin(t) is Annotated:298        return get_load_param_if_needed_function(get_args(t)[0])299 300    def load_base_model(v: Dict[str, Any], t: Type[BaseModel]) -> BaseModel:301        return t(**v)302 303    return load_base_model if isinstance(t, type) and issubclass(t, BaseModel) else None304 305 306def load_basemodels_if_needed(func: Callable[..., Any]) -> Callable[..., Any]:307    """A decorator to load the parameters of a function if they are Pydantic models308 309    Args:310        func: The function with annotated parameters311 312    Returns:313        A function that loads the parameters before calling the original function314 315    """316    # get the type annotations of the parameters317    typed_signature = get_typed_signature(func)318    param_annotations = get_param_annotations(typed_signature)319 320    # get functions for loading BaseModels when needed based on the type annotations321    kwargs_mapping = {k: get_load_param_if_needed_function(t) for k, t in param_annotations.items()}322 323    # remove the None values324    kwargs_mapping = {k: f for k, f in kwargs_mapping.items() if f is not None}325 326    # a function that loads the parameters before calling the original function327    @functools.wraps(func)328    def _load_parameters_if_needed(*args: Any, **kwargs: Any) -> Any:329        # load the BaseModels if needed330        for k, f in kwargs_mapping.items():331            kwargs[k] = f(kwargs[k], param_annotations[k])332 333        # call the original function334        return func(*args, **kwargs)335 336    @functools.wraps(func)337    async def _a_load_parameters_if_needed(*args: Any, **kwargs: Any) -> Any:338        # load the BaseModels if needed339        for k, f in kwargs_mapping.items():340            kwargs[k] = f(kwargs[k], param_annotations[k])341 342        # call the original function343        return await func(*args, **kwargs)344 345    if inspect.iscoroutinefunction(func):346        return _a_load_parameters_if_needed347    else:348        return _load_parameters_if_needed349 350 351def serialize_to_str(x: Any) -> str:352    if isinstance(x, str):353        return x354    elif isinstance(x, BaseModel):355        return model_dump_json(x)356    else:357        return json.dumps(x)358