llama-index

1
from inspect import signature
2
from typing import Any, Callable, List, Optional, Tuple, Type, Union, cast
3

4
from llama_index.legacy.bridge.pydantic import BaseModel, FieldInfo, create_model
5

6

7
def create_schema_from_function(
8
    name: str,
9
    func: Callable[..., Any],
10
    additional_fields: Optional[
11
        List[Union[Tuple[str, Type, Any], Tuple[str, Type]]]
12
    ] = None,
13
) -> Type[BaseModel]:
14
    """Create schema from function."""
15
    fields = {}
16
    params = signature(func).parameters
17
    for param_name in params:
18
        param_type = params[param_name].annotation
19
        param_default = params[param_name].default
20

21
        if param_type is params[param_name].empty:
22
            param_type = Any
23

24
        if param_default is params[param_name].empty:
25
            # Required field
26
            fields[param_name] = (param_type, FieldInfo())
27
        elif isinstance(param_default, FieldInfo):
28
            # Field with pydantic.Field as default value
29
            fields[param_name] = (param_type, param_default)
30
        else:
31
            fields[param_name] = (param_type, FieldInfo(default=param_default))
32

33
    additional_fields = additional_fields or []
34
    for field_info in additional_fields:
35
        if len(field_info) == 3:
36
            field_info = cast(Tuple[str, Type, Any], field_info)
37
            field_name, field_type, field_default = field_info
38
            fields[field_name] = (field_type, FieldInfo(default=field_default))
39
        elif len(field_info) == 2:
40
            # Required field has no default value
41
            field_info = cast(Tuple[str, Type], field_info)
42
            field_name, field_type = field_info
43
            fields[field_name] = (field_type, FieldInfo())
44
        else:
45
            raise ValueError(
46
                f"Invalid additional field info: {field_info}. "
47
                "Must be a tuple of length 2 or 3."
48
            )
49

50
    return create_model(name, **fields)  # type: ignore
51