codekingpro/portable-devtools

1"""Utilities for pydantic."""2 3from __future__ import annotations4 5import inspect6import textwrap7import warnings8from contextlib import nullcontext9from functools import lru_cache, wraps10from types import GenericAlias11from typing import (12    TYPE_CHECKING,13    Any,14    TypeVar,15    cast,16    overload,17)18 19import pydantic20from packaging import version21 22# root_validator is deprecated but we need it for backward compatibility of @pre_init23from pydantic import (  # type: ignore[deprecated]24    BaseModel,25    ConfigDict,26    Field,27    PydanticDeprecationWarning,28    RootModel,29    root_validator,30)31from pydantic import (32    create_model as _create_model_base,33)34from pydantic.fields import FieldInfo as FieldInfoV235from pydantic.json_schema import (36    DEFAULT_REF_TEMPLATE,37    GenerateJsonSchema,38    JsonSchemaMode,39    JsonSchemaValue,40)41from pydantic.v1 import BaseModel as BaseModelV142from pydantic.v1 import create_model as create_model_v143from typing_extensions import deprecated, override44 45if TYPE_CHECKING:46    from collections.abc import Callable47 48    from pydantic.v1.fields import ModelField49    from pydantic_core import core_schema50 51PYDANTIC_VERSION = version.parse(pydantic.__version__)52 53 54@deprecated("Use PYDANTIC_VERSION.major instead.")55def get_pydantic_major_version() -> int:56    """DEPRECATED - Get the major version of Pydantic.57 58    Use `PYDANTIC_VERSION.major` instead.59 60    Returns:61        The major version of Pydantic.62    """63    return PYDANTIC_VERSION.major64 65 66PYDANTIC_MAJOR_VERSION = PYDANTIC_VERSION.major67PYDANTIC_MINOR_VERSION = PYDANTIC_VERSION.minor68 69IS_PYDANTIC_V1 = False70IS_PYDANTIC_V2 = True71 72PydanticBaseModel = BaseModel73TypeBaseModel = type[BaseModel]74 75TBaseModel = TypeVar("TBaseModel", bound=PydanticBaseModel)76 77 78def is_pydantic_v1_subclass(cls: type) -> bool:79    """Check if the given class is Pydantic v1-like.80 81    Returns:82        `True` if the given class is a subclass of Pydantic `BaseModel` 1.x.83    """84    return issubclass(cls, BaseModelV1)85 86 87def is_pydantic_v2_subclass(cls: type) -> bool:88    """Check if the given class is Pydantic v2-like.89 90    Returns:91        `True` if the given class is a subclass of Pydantic `BaseModel` 2.x.92    """93    return issubclass(cls, BaseModel)94 95 96def is_basemodel_subclass(cls: type) -> bool:97    """Check if the given class is a subclass of Pydantic `BaseModel`.98 99    Check if the given class is a subclass of any of the following:100 101    * `pydantic.BaseModel` in Pydantic 2.x102    * `pydantic.v1.BaseModel` in Pydantic 2.x103 104    Returns:105        `True` if the given class is a subclass of Pydantic `BaseModel`.106    """107    # Before we can use issubclass on the cls we need to check if it is a class108    if not inspect.isclass(cls) or isinstance(cls, GenericAlias):109        return False110 111    return issubclass(cls, (BaseModel, BaseModelV1))112 113 114def is_basemodel_instance(obj: Any) -> bool:115    """Check if the given class is an instance of Pydantic `BaseModel`.116 117    Check if the given class is an instance of any of the following:118 119    * `pydantic.BaseModel` in Pydantic 2.x120    * `pydantic.v1.BaseModel` in Pydantic 2.x121 122    Returns:123        `True` if the given class is an instance of Pydantic `BaseModel`.124    """125    return isinstance(obj, (BaseModel, BaseModelV1))126 127 128# How to type hint this?129def pre_init(func: Callable) -> Any:130    """Decorator to run a function before model initialization.131 132    Args:133        func: The function to run before model initialization.134 135    Returns:136        The decorated function.137    """138    with warnings.catch_warnings():139        warnings.filterwarnings(action="ignore", category=PydanticDeprecationWarning)140 141        # Ideally we would use @model_validator(mode="before") but this would change the142        # order of the validators. See https://github.com/pydantic/pydantic/discussions/7434.143        # So we keep root_validator for backward compatibility.144        @root_validator(pre=True)  # type: ignore[deprecated]145        @wraps(func)146        def wrapper(cls: type[BaseModel], values: dict[str, Any]) -> Any:147            """Decorator to run a function before model initialization.148 149            Args:150                cls: The model class.151                values: The values to initialize the model with.152 153            Returns:154                The values to initialize the model with.155            """156            # Insert default values157            fields = cls.model_fields158            for name, field_info in fields.items():159                # Check if allow_population_by_field_name is enabled160                # If yes, then set the field name to the alias161                if (162                    hasattr(cls, "Config")163                    and hasattr(cls.Config, "allow_population_by_field_name")164                    and cls.Config.allow_population_by_field_name165                    and field_info.alias in values166                ):167                    values[name] = values.pop(field_info.alias)168                if (169                    hasattr(cls, "model_config")170                    and cls.model_config.get("populate_by_name")171                    and field_info.alias in values172                ):173                    values[name] = values.pop(field_info.alias)174 175                if (176                    name not in values or values[name] is None177                ) and not field_info.is_required():178                    if field_info.default_factory is not None:179                        values[name] = field_info.default_factory()  # type: ignore[call-arg]180                    else:181                        values[name] = field_info.default182 183            # Call the decorated function184            return func(cls, values)185 186    return wrapper187 188 189class _IgnoreUnserializable(GenerateJsonSchema):190    """A JSON schema generator that ignores unknown types.191 192    https://docs.pydantic.dev/latest/concepts/json_schema/#customizing-the-json-schema-generation-process193    """194 195    @override196    def handle_invalid_for_json_schema(197        self, schema: core_schema.CoreSchema, error_info: str198    ) -> JsonSchemaValue:199        return {}200 201 202def _create_subset_model_v1(203    name: str,204    model: type[BaseModelV1],205    field_names: list,206    *,207    descriptions: dict | None = None,208    fn_description: str | None = None,209) -> type[BaseModelV1]:210    """Create a Pydantic model with only a subset of model's fields."""211    fields = {}212 213    for field_name in field_names:214        # Using pydantic v1 so can access __fields__ as a dict.215        field = model.__fields__[field_name]216        t = (217            # this isn't perfect but should work for most functions218            field.outer_type_219            if field.required and not field.allow_none220            else field.outer_type_ | None221        )222        if descriptions and field_name in descriptions:223            field.field_info.description = descriptions[field_name]224        fields[field_name] = (t, field.field_info)225 226    rtn = cast("type[BaseModelV1]", create_model_v1(name, **fields))  # type: ignore[call-overload]227    rtn.__doc__ = textwrap.dedent(fn_description or model.__doc__ or "")228    return rtn229 230 231def _create_subset_model_v2(232    name: str,233    model: type[BaseModel],234    field_names: list[str],235    *,236    descriptions: dict | None = None,237    fn_description: str | None = None,238) -> type[BaseModel]:239    """Create a Pydantic model with a subset of the model fields."""240    descriptions_ = descriptions or {}241    fields = {}242    for field_name in field_names:243        field = model.model_fields[field_name]244        description = descriptions_.get(field_name, field.description)245        field_kwargs: dict[str, Any] = {"description": description}246        if field.default_factory is not None:247            field_kwargs["default_factory"] = field.default_factory248        else:249            field_kwargs["default"] = field.default250        field_info = FieldInfoV2(**field_kwargs)251        if field.metadata:252            field_info.metadata = field.metadata253        fields[field_name] = (field.annotation, field_info)254 255    rtn = cast(256        "type[BaseModel]",257        _create_model_base(  # type: ignore[call-overload]258            name, **fields, __config__=ConfigDict(arbitrary_types_allowed=True)259        ),260    )261 262    # TODO(0.3): Determine if there is a more "pydantic" way to preserve annotations.263    # This is done to preserve __annotations__ when working with pydantic 2.x264    # and using the Annotated type with TypedDict.265    # Comment out the following line, to trigger the relevant test case.266    selected_annotations = [267        (name, annotation)268        for name, annotation in model.__annotations__.items()269        if name in field_names270    ]271 272    rtn.__annotations__ = dict(selected_annotations)273    rtn.__doc__ = textwrap.dedent(fn_description or model.__doc__ or "")274    return rtn275 276 277# Private functionality to create a subset model that's compatible across278# different versions of pydantic.279# Handles pydantic versions 2.x. including v1 of pydantic in 2.x.280# However, can't find a way to type hint this.281def _create_subset_model(282    name: str,283    model: TypeBaseModel,284    field_names: list[str],285    *,286    descriptions: dict | None = None,287    fn_description: str | None = None,288) -> type[BaseModel]:289    """Create subset model using the same pydantic version as the input model.290 291    Returns:292        The created subset model.293    """294    if issubclass(model, BaseModelV1):295        return _create_subset_model_v1(296            name,297            model,298            field_names,299            descriptions=descriptions,300            fn_description=fn_description,301        )302    return _create_subset_model_v2(303        name,304        model,305        field_names,306        descriptions=descriptions,307        fn_description=fn_description,308    )309 310 311@overload312def get_fields(model: type[BaseModel]) -> dict[str, FieldInfoV2]: ...313 314 315@overload316def get_fields(model: BaseModel) -> dict[str, FieldInfoV2]: ...317 318 319@overload320def get_fields(model: type[BaseModelV1]) -> dict[str, ModelField]: ...321 322 323@overload324def get_fields(model: BaseModelV1) -> dict[str, ModelField]: ...325 326 327def get_fields(328    model: type[BaseModel | BaseModelV1] | BaseModel | BaseModelV1,329) -> dict[str, FieldInfoV2] | dict[str, ModelField]:330    """Return the field names of a Pydantic model.331 332    Args:333        model: The Pydantic model or instance.334 335    Raises:336        TypeError: If the model is not a Pydantic model.337    """338    if not isinstance(model, type):339        model = type(model)340    if issubclass(model, BaseModel):341        return model.model_fields342    if issubclass(model, BaseModelV1):343        return model.__fields__344    msg = f"Expected a Pydantic model. Got {model}"345    raise TypeError(msg)346 347 348_SchemaConfig = ConfigDict(349    arbitrary_types_allowed=True, frozen=True, protected_namespaces=()350)351 352NO_DEFAULT = object()353 354 355def _create_root_model(356    name: str,357    type_: Any,358    module_name: str | None = None,359    default_: object = NO_DEFAULT,360) -> type[BaseModel]:361    """Create a base class."""362 363    def schema(364        cls: type[BaseModelV1],365        by_alias: bool = True,  # noqa: FBT001,FBT002366        ref_template: str = DEFAULT_REF_TEMPLATE,367    ) -> dict[str, Any]:368        super_cls = cast("type[BaseModelV1]", super(cls, cls))369        schema_ = super_cls.schema(by_alias=by_alias, ref_template=ref_template)370        schema_["title"] = name371        return schema_372 373    def model_json_schema(374        cls: type[BaseModel],375        by_alias: bool = True,  # noqa: FBT001,FBT002376        ref_template: str = DEFAULT_REF_TEMPLATE,377        schema_generator: type[GenerateJsonSchema] = GenerateJsonSchema,378        mode: JsonSchemaMode = "validation",379    ) -> dict[str, Any]:380        super_cls = cast("type[BaseModel]", super(cls, cls))381        schema_ = super_cls.model_json_schema(382            by_alias=by_alias,383            ref_template=ref_template,384            schema_generator=schema_generator,385            mode=mode,386        )387        schema_["title"] = name388        return schema_389 390    base_class_attributes = {391        "__annotations__": {"root": type_},392        "model_config": ConfigDict(arbitrary_types_allowed=True),393        "schema": classmethod(schema),394        "model_json_schema": classmethod(model_json_schema),395        "__module__": module_name or "langchain_core.runnables.utils",396    }397 398    if default_ is not NO_DEFAULT:399        base_class_attributes["root"] = default_400    with warnings.catch_warnings():401        try:402            if (403                isinstance(type_, type)404                and not isinstance(type_, GenericAlias)405                and issubclass(type_, BaseModelV1)406            ):407                warnings.filterwarnings(408                    action="ignore", category=PydanticDeprecationWarning409                )410        except TypeError:411            pass412        custom_root_type = type(name, (RootModel,), base_class_attributes)413    return cast("type[BaseModel]", custom_root_type)414 415 416@lru_cache(maxsize=256)417def _create_root_model_cached(418    model_name: str,419    type_: Any,420    *,421    module_name: str | None = None,422    default_: object = NO_DEFAULT,423) -> type[BaseModel]:424    return _create_root_model(425        model_name, type_, default_=default_, module_name=module_name426    )427 428 429@lru_cache(maxsize=256)430def _create_model_cached(431    model_name: str,432    /,433    **field_definitions: Any,434) -> type[BaseModel]:435    return _create_model_base(436        model_name,437        __config__=_SchemaConfig,438        **_remap_field_definitions(field_definitions),439    )440 441 442def create_model(443    model_name: str,444    module_name: str | None = None,445    /,446    **field_definitions: Any,447) -> type[BaseModel]:448    """Create a Pydantic model with the given field definitions.449 450    Please use `create_model_v2` instead of this function.451 452    Args:453        model_name: The name of the model.454        module_name: The name of the module where the model is defined.455 456            This is used by Pydantic to resolve any forward references.457        **field_definitions: The field definitions for the model.458 459    Returns:460        The created model.461    """462    kwargs = {}463    if "__root__" in field_definitions:464        kwargs["root"] = field_definitions.pop("__root__")465 466    return create_model_v2(467        model_name,468        module_name=module_name,469        field_definitions=field_definitions,470        **kwargs,471    )472 473 474# Reserved names should capture all the `public` names / methods that are475# used by BaseModel internally. This will keep the reserved names up-to-date.476# For reference, the reserved names are:477# "construct", "copy", "dict", "from_orm", "json", "parse_file", "parse_obj",478# "parse_raw", "schema", "schema_json", "update_forward_refs", "validate",479# "model_computed_fields", "model_config", "model_construct", "model_copy",480# "model_dump", "model_dump_json", "model_extra", "model_fields",481# "model_fields_set", "model_json_schema", "model_parametrized_name",482# "model_post_init", "model_rebuild", "model_validate", "model_validate_json",483# "model_validate_strings"484_RESERVED_NAMES = {key for key in dir(BaseModel) if not key.startswith("_")}485 486 487def _remap_field_definitions(field_definitions: dict[str, Any]) -> dict[str, Any]:488    """This remaps fields to avoid colliding with internal pydantic fields."""489    remapped = {}490    for key, value in field_definitions.items():491        if key.startswith("_") or key in _RESERVED_NAMES:492            # Let's add a prefix to avoid colliding with internal pydantic fields493            if isinstance(value, FieldInfoV2):494                msg = (495                    f"Remapping for fields starting with '_' or fields with a name "496                    f"matching a reserved name {_RESERVED_NAMES} is not supported if "497                    f" the field is a pydantic Field instance. Got {key}."498                )499                raise NotImplementedError(msg)500            type_, default_ = value501            remapped[f"private_{key}"] = (502                type_,503                Field(504                    default=default_,505                    alias=key,506                    serialization_alias=key,507                    title=key.lstrip("_").replace("_", " ").title(),508                ),509            )510        else:511            remapped[key] = value512    return remapped513 514 515def create_model_v2(516    model_name: str,517    *,518    module_name: str | None = None,519    field_definitions: dict[str, Any] | None = None,520    root: Any | None = None,521) -> type[BaseModel]:522    """Create a Pydantic model with the given field definitions.523 524    !!! warning525 526        Do not use outside of langchain packages. This API is subject to change at any527        time.528 529    Args:530        model_name: The name of the model.531        module_name: The name of the module where the model is defined.532 533            This is used by Pydantic to resolve any forward references.534        field_definitions: The field definitions for the model.535        root: Type for a root model (`RootModel`)536 537    Returns:538        The created model.539    """540    field_definitions = field_definitions or {}541 542    if root:543        if field_definitions:544            msg = (545                "When specifying __root__ no other "546                f"fields should be provided. Got {field_definitions}"547            )548            raise NotImplementedError(msg)549 550        if isinstance(root, tuple):551            kwargs = {"type_": root[0], "default_": root[1]}552        else:553            kwargs = {"type_": root}554 555        try:556            named_root_model = _create_root_model_cached(557                model_name, module_name=module_name, **kwargs558            )559        except TypeError:560            # something in the arguments into _create_root_model_cached is not hashable561            named_root_model = _create_root_model(562                model_name,563                module_name=module_name,564                **kwargs,565            )566        return named_root_model567 568    # No root, just field definitions569    names = set(field_definitions.keys())570 571    capture_warnings = False572 573    for name in names:574        # Also if any non-reserved name is used (e.g., model_id or model_name)575        if name.startswith("model"):576            capture_warnings = True577 578    with warnings.catch_warnings() if capture_warnings else nullcontext():579        if capture_warnings:580            warnings.filterwarnings(action="ignore")581        try:582            return _create_model_cached(model_name, **field_definitions)583        except TypeError:584            # something in field definitions is not hashable585            return _create_model_base(586                model_name,587                __config__=_SchemaConfig,588                **_remap_field_definitions(field_definitions),589            )590