1"""Utilities for pydantic."""23from __future__ import annotations45import 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)1819import pydantic20from packaging import version2122# 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, override4445if TYPE_CHECKING:46    from collections.abc import Callable4748    from pydantic.v1.fields import ModelField49    from pydantic_core import core_schema5051PYDANTIC_VERSION = version.parse(pydantic.__version__)525354@deprecated("Use PYDANTIC_VERSION.major instead.")55def get_pydantic_major_version() -> int:56    """DEPRECATED - Get the major version of Pydantic.5758    Use `PYDANTIC_VERSION.major` instead.5960    Returns:61        The major version of Pydantic.62    """63    return PYDANTIC_VERSION.major646566PYDANTIC_MAJOR_VERSION = PYDANTIC_VERSION.major67PYDANTIC_MINOR_VERSION = PYDANTIC_VERSION.minor6869IS_PYDANTIC_V1 = False70IS_PYDANTIC_V2 = True7172PydanticBaseModel = BaseModel | BaseModelV173TypeBaseModel = type[BaseModel] | type[BaseModelV1]7475TBaseModel = TypeVar("TBaseModel", bound=PydanticBaseModel)767778def is_pydantic_v1_subclass(cls: type) -> bool:79    """Check if the given class is Pydantic v1-like.8081    Returns:82        `True` if the given class is a subclass of Pydantic `BaseModel` 1.x.83    """84    return issubclass(cls, BaseModelV1)858687def is_pydantic_v2_subclass(cls: type) -> bool:88    """Check if the given class is Pydantic v2-like.8990    Returns:91        `True` if the given class is a subclass of Pydantic `BaseModel` 2.x.92    """93    return issubclass(cls, BaseModel)949596def is_basemodel_subclass(cls: type) -> bool:97    """Check if the given class is a subclass of Pydantic `BaseModel`.9899    Check if the given class is a subclass of any of the following:100101    * `pydantic.BaseModel` in Pydantic 2.x102    * `pydantic.v1.BaseModel` in Pydantic 2.x103104    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 False110111    return issubclass(cls, (BaseModel, BaseModelV1))112113114def is_basemodel_instance(obj: Any) -> bool:115    """Check if the given class is an instance of Pydantic `BaseModel`.116117    Check if the given class is an instance of any of the following:118119    * `pydantic.BaseModel` in Pydantic 2.x120    * `pydantic.v1.BaseModel` in Pydantic 2.x121122    Returns:123        `True` if the given class is an instance of Pydantic `BaseModel`.124    """125    return isinstance(obj, (BaseModel, BaseModelV1))126127128# How to type hint this?129def pre_init(130    func: Callable[[Any, dict[str, Any]], Any],131) -> Callable[[Any, dict[str, Any]], Any]:132    """Decorator to run a function before model initialization.133134    Args:135        func: The function to run before model initialization.136137    Returns:138        The decorated function.139    """140    with warnings.catch_warnings():141        warnings.filterwarnings(action="ignore", category=PydanticDeprecationWarning)142143        # Ideally we would use @model_validator(mode="before") but this would change the144        # order of the validators. See https://github.com/pydantic/pydantic/discussions/7434.145        # So we keep root_validator for backward compatibility.146        @root_validator(pre=True)  # type: ignore[deprecated]147        @wraps(func)148        def wrapper(cls: type[BaseModel], values: dict[str, Any]) -> Any:149            """Decorator to run a function before model initialization.150151            Args:152                cls: The model class.153                values: The values to initialize the model with.154155            Returns:156                The values to initialize the model with.157            """158            # Insert default values159            fields = cls.model_fields160            for name, field_info in fields.items():161                # Check if allow_population_by_field_name is enabled162                # If yes, then set the field name to the alias163                if (164                    hasattr(cls, "Config")165                    and hasattr(cls.Config, "allow_population_by_field_name")166                    and cls.Config.allow_population_by_field_name167                    and field_info.alias in values168                ):169                    values[name] = values.pop(field_info.alias)170                if (171                    hasattr(cls, "model_config")172                    and cls.model_config.get("populate_by_name")173                    and field_info.alias in values174                ):175                    values[name] = values.pop(field_info.alias)176177                if (178                    name not in values or values[name] is None179                ) and not field_info.is_required():180                    if field_info.default_factory is not None:181                        values[name] = field_info.default_factory()  # type: ignore[call-arg]182                    else:183                        values[name] = field_info.default184185            # Call the decorated function186            return func(cls, values)187188    return wrapper189190191class _IgnoreUnserializable(GenerateJsonSchema):192    """A JSON schema generator that ignores unknown types.193194    https://docs.pydantic.dev/latest/concepts/json_schema/#customizing-the-json-schema-generation-process195    """196197    @override198    def handle_invalid_for_json_schema(199        self, schema: core_schema.CoreSchema, error_info: str200    ) -> JsonSchemaValue:201        return {}202203204def _create_subset_model_v1(205    name: str,206    model: type[BaseModelV1],207    field_names: list[str],208    *,209    descriptions: dict[str, str] | None = None,210    fn_description: str | None = None,211) -> type[BaseModelV1]:212    """Create a Pydantic model with only a subset of model's fields."""213    fields = {}214215    for field_name in field_names:216        # Using pydantic v1 so can access __fields__ as a dict.217        field = model.__fields__[field_name]218        t = (219            # this isn't perfect but should work for most functions220            field.outer_type_221            if field.required and not field.allow_none222            else field.outer_type_ | None223        )224        if descriptions and field_name in descriptions:225            field.field_info.description = descriptions[field_name]226        fields[field_name] = (t, field.field_info)227228    rtn = cast("type[BaseModelV1]", create_model_v1(name, **fields))  # type: ignore[call-overload]229    rtn.__doc__ = textwrap.dedent(fn_description or model.__doc__ or "")230    return rtn231232233def _create_subset_model_v2(234    name: str,235    model: type[BaseModel],236    field_names: list[str],237    *,238    descriptions: dict[str, str] | None = None,239    fn_description: str | None = None,240) -> type[BaseModel]:241    """Create a Pydantic model with a subset of the model fields."""242    descriptions_ = descriptions or {}243    fields = {}244    for field_name in field_names:245        field = model.model_fields[field_name]246        description = descriptions_.get(field_name, field.description)247        field_kwargs: dict[str, Any] = {"description": description}248        if field.default_factory is not None:249            field_kwargs["default_factory"] = field.default_factory250        else:251            field_kwargs["default"] = field.default252        field_info = FieldInfoV2(**field_kwargs)253        if field.metadata:254            field_info.metadata = field.metadata255        fields[field_name] = (field.annotation, field_info)256257    rtn = cast(258        "type[BaseModel]",259        _create_model_base(  # type: ignore[call-overload]260            name, **fields, __config__=ConfigDict(arbitrary_types_allowed=True)261        ),262    )263264    # TODO(0.3): Determine if there is a more "pydantic" way to preserve annotations.265    # This is done to preserve __annotations__ when working with pydantic 2.x266    # and using the Annotated type with TypedDict.267    # Comment out the following line, to trigger the relevant test case.268    selected_annotations = [269        (name, annotation)270        for name, annotation in model.__annotations__.items()271        if name in field_names272    ]273274    rtn.__annotations__ = dict(selected_annotations)275    rtn.__doc__ = textwrap.dedent(fn_description or model.__doc__ or "")276    return rtn277278279# Private functionality to create a subset model that's compatible across280# different versions of pydantic.281# Handles pydantic versions 2.x. including v1 of pydantic in 2.x.282# However, can't find a way to type hint this.283def _create_subset_model(284    name: str,285    model: TypeBaseModel,286    field_names: list[str],287    *,288    descriptions: dict[str, str] | None = None,289    fn_description: str | None = None,290) -> TypeBaseModel:291    """Create subset model using the same pydantic version as the input model.292293    Returns:294        The created subset model.295    """296    if issubclass(model, BaseModelV1):297        return _create_subset_model_v1(298            name,299            model,300            field_names,301            descriptions=descriptions,302            fn_description=fn_description,303        )304    return _create_subset_model_v2(305        name,306        model,307        field_names,308        descriptions=descriptions,309        fn_description=fn_description,310    )311312313@overload314def get_fields(model: type[BaseModel]) -> dict[str, FieldInfoV2]: ...315316317@overload318def get_fields(model: BaseModel) -> dict[str, FieldInfoV2]: ...319320321@overload322def get_fields(model: type[BaseModelV1]) -> dict[str, ModelField]: ...323324325@overload326def get_fields(model: BaseModelV1) -> dict[str, ModelField]: ...327328329def get_fields(330    model: type[BaseModel | BaseModelV1] | BaseModel | BaseModelV1,331) -> dict[str, FieldInfoV2] | dict[str, ModelField]:332    """Return the field names of a Pydantic model.333334    Args:335        model: The Pydantic model or instance.336337    Raises:338        TypeError: If the model is not a Pydantic model.339    """340    if not isinstance(model, type):341        model = type(model)342    if issubclass(model, BaseModel):343        return model.model_fields344    if issubclass(model, BaseModelV1):345        return model.__fields__346    msg = f"Expected a Pydantic model. Got {model}"  # type: ignore[unreachable]347    raise TypeError(msg)348349350def model_json_schema(model: TypeBaseModel) -> dict[str, Any]:351    """Return the JSON schema of a Pydantic model class of either major version.352353    Dispatches to the correct method for Pydantic v1 (`schema`) or v2354    (`model_json_schema`), so callers holding a `TypeBaseModel` don't have to355    branch on the model's version themselves.356357    Args:358        model: The Pydantic model class.359360    Raises:361        TypeError: If the model is not a Pydantic model class.362    """363    if issubclass(model, BaseModel):364        return model.model_json_schema()365    if issubclass(model, BaseModelV1):366        return model.schema()367    msg = f"Expected a Pydantic model. Got {model}"  # type: ignore[unreachable]368    raise TypeError(msg)369370371def model_validate(model: TypeBaseModel, obj: Any) -> PydanticBaseModel:372    """Validate `obj` against a Pydantic model class of either major version.373374    Dispatches to the correct method for Pydantic v1 (`parse_obj`) or v2375    (`model_validate`), so callers holding a `TypeBaseModel` don't have to376    branch on the model's version themselves.377378    Args:379        model: The Pydantic model class to validate against.380        obj: The object to validate.381382    Raises:383        TypeError: If the model is not a Pydantic model class.384    """385    if issubclass(model, BaseModel):386        return model.model_validate(obj)387    if issubclass(model, BaseModelV1):388        return model.parse_obj(obj)389    msg = f"Expected a Pydantic model. Got {model}"  # type: ignore[unreachable]390    raise TypeError(msg)391392393_SchemaConfig = ConfigDict(394    arbitrary_types_allowed=True, frozen=True, protected_namespaces=()395)396397NO_DEFAULT = object()398399400def _create_root_model(401    name: str,402    type_: Any,403    module_name: str | None = None,404    default_: object = NO_DEFAULT,405) -> type[BaseModel]:406    """Create a base class."""407408    def schema(409        cls: type[BaseModelV1],410        by_alias: bool = True,  # noqa: FBT001,FBT002411        ref_template: str = DEFAULT_REF_TEMPLATE,412    ) -> dict[str, Any]:413        super_cls = cast("type[BaseModelV1]", super(cls, cls))414        schema_ = super_cls.schema(by_alias=by_alias, ref_template=ref_template)415        schema_["title"] = name416        return schema_417418    def model_json_schema(419        cls: type[BaseModel],420        by_alias: bool = True,  # noqa: FBT001,FBT002421        ref_template: str = DEFAULT_REF_TEMPLATE,422        schema_generator: type[GenerateJsonSchema] = GenerateJsonSchema,423        mode: JsonSchemaMode = "validation",424    ) -> dict[str, Any]:425        super_cls = cast("type[BaseModel]", super(cls, cls))426        schema_ = super_cls.model_json_schema(427            by_alias=by_alias,428            ref_template=ref_template,429            schema_generator=schema_generator,430            mode=mode,431        )432        schema_["title"] = name433        return schema_434435    base_class_attributes = {436        "__annotations__": {"root": type_},437        "model_config": ConfigDict(arbitrary_types_allowed=True),438        "schema": classmethod(schema),439        "model_json_schema": classmethod(model_json_schema),440        "__module__": module_name or "langchain_core.runnables.utils",441    }442443    if default_ is not NO_DEFAULT:444        base_class_attributes["root"] = default_445    with warnings.catch_warnings():446        try:447            if isinstance(type_, type) and issubclass(type_, BaseModelV1):448                warnings.filterwarnings(449                    action="ignore", category=PydanticDeprecationWarning450                )451        except TypeError:452            pass453        custom_root_type = type(name, (RootModel,), base_class_attributes)454    return cast("type[BaseModel]", custom_root_type)455456457@lru_cache(maxsize=256)458def _create_root_model_cached(459    model_name: str,460    type_: Any,461    *,462    module_name: str | None = None,463    default_: object = NO_DEFAULT,464) -> type[BaseModel]:465    return _create_root_model(466        model_name, type_, default_=default_, module_name=module_name467    )468469470@lru_cache(maxsize=256)471def _create_model_cached(472    model_name: str,473    /,474    **field_definitions: Any,475) -> type[BaseModel]:476    return _create_model_base(477        model_name,478        __config__=_SchemaConfig,479        **_remap_field_definitions(field_definitions),480    )481482483def create_model(484    model_name: str,485    module_name: str | None = None,486    /,487    **field_definitions: Any,488) -> type[BaseModel]:489    """Create a Pydantic model with the given field definitions.490491    Please use `create_model_v2` instead of this function.492493    Args:494        model_name: The name of the model.495        module_name: The name of the module where the model is defined.496497            This is used by Pydantic to resolve any forward references.498        **field_definitions: The field definitions for the model.499500    Returns:501        The created model.502    """503    kwargs = {}504    if "__root__" in field_definitions:505        kwargs["root"] = field_definitions.pop("__root__")506507    return create_model_v2(508        model_name,509        module_name=module_name,510        field_definitions=field_definitions,511        **kwargs,512    )513514515# Reserved names should capture all the `public` names / methods that are516# used by BaseModel internally. This will keep the reserved names up-to-date.517# For reference, the reserved names are:518# "construct", "copy", "dict", "from_orm", "json", "parse_file", "parse_obj",519# "parse_raw", "schema", "schema_json", "update_forward_refs", "validate",520# "model_computed_fields", "model_config", "model_construct", "model_copy",521# "model_dump", "model_dump_json", "model_extra", "model_fields",522# "model_fields_set", "model_json_schema", "model_parametrized_name",523# "model_post_init", "model_rebuild", "model_validate", "model_validate_json",524# "model_validate_strings"525_RESERVED_NAMES = {key for key in dir(BaseModel) if not key.startswith("_")}526527528def _remap_field_definitions(field_definitions: dict[str, Any]) -> dict[str, Any]:529    """This remaps fields to avoid colliding with internal pydantic fields."""530    remapped = {}531    for key, value in field_definitions.items():532        if key.startswith("_") or key in _RESERVED_NAMES:533            # Let's add a prefix to avoid colliding with internal pydantic fields534            if isinstance(value, FieldInfoV2):535                msg = (536                    f"Remapping for fields starting with '_' or fields with a name "537                    f"matching a reserved name {_RESERVED_NAMES} is not supported if "538                    f" the field is a pydantic Field instance. Got {key}."539                )540                raise NotImplementedError(msg)541            type_, default_ = value542            remapped[f"private_{key}"] = (543                type_,544                Field(545                    default=default_,546                    alias=key,547                    serialization_alias=key,548                    title=key.lstrip("_").replace("_", " ").title(),549                ),550            )551        else:552            remapped[key] = value553    return remapped554555556def create_model_v2(557    model_name: str,558    *,559    module_name: str | None = None,560    field_definitions: dict[str, Any] | None = None,561    root: Any | None = None,562) -> type[BaseModel]:563    """Create a Pydantic model with the given field definitions.564565    !!! warning566567        Do not use outside of langchain packages. This API is subject to change at any568        time.569570    Args:571        model_name: The name of the model.572        module_name: The name of the module where the model is defined.573574            This is used by Pydantic to resolve any forward references.575        field_definitions: The field definitions for the model.576        root: Type for a root model (`RootModel`)577578    Returns:579        The created model.580    """581    field_definitions = field_definitions or {}582583    if root:584        if field_definitions:585            msg = (586                "When specifying __root__ no other "587                f"fields should be provided. Got {field_definitions}"588            )589            raise NotImplementedError(msg)590591        if isinstance(root, tuple):592            kwargs = {"type_": root[0], "default_": root[1]}593        else:594            kwargs = {"type_": root}595596        try:597            named_root_model = _create_root_model_cached(598                model_name, module_name=module_name, **kwargs599            )600        except TypeError:601            # something in the arguments into _create_root_model_cached is not hashable602            named_root_model = _create_root_model(603                model_name,604                module_name=module_name,605                **kwargs,606            )607        return named_root_model608609    # No root, just field definitions610    names = set(field_definitions.keys())611612    capture_warnings = False613614    for name in names:615        # Also if any non-reserved name is used (e.g., model_id or model_name)616        if name.startswith("model"):617            capture_warnings = True618619    with warnings.catch_warnings() if capture_warnings else nullcontext():620        if capture_warnings:621            warnings.filterwarnings(action="ignore")622        try:623            return _create_model_cached(model_name, **field_definitions)624        except TypeError:625            # something in field definitions is not hashable626            return _create_model_base(627                model_name,628                __config__=_SchemaConfig,629                **_remap_field_definitions(field_definitions),630            )