Cleanup Pydantic models and handle deprecation warnings

cbornet · cbornet · commit 75293669870a · 2025-04-12T03:05:03.000+02:00
diff --git a/libs/core/langchain_core/_api/deprecation.py b/libs/core/langchain_core/_api/deprecation.py
@@ -23,6 +23,8 @@
     cast,
 )
 
+from pydantic.fields import FieldInfo
+from pydantic.v1.fields import FieldInfo as FieldInfoV1
 from typing_extensions import ParamSpec
 
 from langchain_core._api.internal import is_caller_internal
@@ -152,10 +154,6 @@ def deprecate(
         _package: str = package,
     ) -> T:
         """Implementation of the decorator returned by `deprecated`."""
-        from langchain_core.utils.pydantic import (  # type: ignore[attr-defined]
-            FieldInfoV1,
-            FieldInfoV2,
-        )
 
         def emit_warning() -> None:
             """Emit the warning."""
@@ -249,7 +247,7 @@ def finalize(wrapper: Callable[..., Any], new_doc: str) -> T:  # noqa: ARG001
                     ),
                 )
 
-        elif isinstance(obj, FieldInfoV2):
+        elif isinstance(obj, FieldInfo):
             wrapped = None
             if not _obj_type:
                 _obj_type = "attribute"
@@ -261,7 +259,7 @@ def finalize(wrapper: Callable[..., Any], new_doc: str) -> T:  # noqa: ARG001
             def finalize(wrapper: Callable[..., Any], new_doc: str) -> T:  # noqa: ARG001
                 return cast(
                     "T",
-                    FieldInfoV2(
+                    FieldInfo(
                         default=obj.default,
                         default_factory=obj.default_factory,
                         description=new_doc,
diff --git a/libs/core/langchain_core/language_models/chat_models.py b/libs/core/langchain_core/language_models/chat_models.py
@@ -584,7 +584,7 @@ def _get_invocation_params(
         stop: Optional[list[str]] = None,
         **kwargs: Any,
     ) -> dict:
-        params = self.dict()
+        params = self.model_dump()
         params["stop"] = stop
         return {**params, **kwargs}
 
@@ -1245,7 +1245,7 @@ def _llm_type(self) -> str:
         """Return type of chat model."""
 
     @override
-    def dict(self, **kwargs: Any) -> dict:
+    def model_dump(self, **kwargs: Any) -> dict:
         """Return a dictionary of the LLM."""
         starter_dict = dict(self._identifying_params)
         starter_dict["_type"] = self._llm_type
diff --git a/libs/core/langchain_core/language_models/llms.py b/libs/core/langchain_core/language_models/llms.py
@@ -528,7 +528,7 @@ def stream(
         else:
             prompt = self._convert_input(input).to_string()
             config = ensure_config(config)
-            params = self.dict()
+            params = self.model_dump()
             params["stop"] = stop
             params = {**params, **kwargs}
             options = {"stop": stop}
@@ -598,7 +598,7 @@ async def astream(
 
         prompt = self._convert_input(input).to_string()
         config = ensure_config(config)
-        params = self.dict()
+        params = self.model_dump()
         params["stop"] = stop
         params = {**params, **kwargs}
         options = {"stop": stop}
@@ -941,7 +941,7 @@ def generate(
             ] * len(prompts)
             run_name_list = [cast("Optional[str]", run_name)] * len(prompts)
         run_ids_list = self._get_run_ids_list(run_id, prompts)
-        params = self.dict()
+        params = self.model_dump()
         params["stop"] = stop
         options = {"stop": stop}
         (
@@ -1193,7 +1193,7 @@ async def agenerate(
             ] * len(prompts)
             run_name_list = [cast("Optional[str]", run_name)] * len(prompts)
         run_ids_list = self._get_run_ids_list(run_id, prompts)
-        params = self.dict()
+        params = self.model_dump()
         params["stop"] = stop
         options = {"stop": stop}
         (
@@ -1400,7 +1400,7 @@ def _llm_type(self) -> str:
         """Return type of llm."""
 
     @override
-    def dict(self, **kwargs: Any) -> dict:
+    def model_dump(self, **kwargs: Any) -> dict:
         """Return a dictionary of the LLM."""
         starter_dict = dict(self._identifying_params)
         starter_dict["_type"] = self._llm_type
@@ -1427,7 +1427,7 @@ def save(self, file_path: Union[Path, str]) -> None:
         directory_path.mkdir(parents=True, exist_ok=True)
 
         # Fetch dictionary to save
-        prompt_dict = self.dict()
+        prompt_dict = self.model_dump()
 
         if save_path.suffix == ".json":
             with save_path.open("w") as f:
diff --git a/libs/core/langchain_core/output_parsers/base.py b/libs/core/langchain_core/output_parsers/base.py
@@ -324,9 +324,9 @@ def _type(self) -> str:
         )
         raise NotImplementedError(msg)
 
-    def dict(self, **kwargs: Any) -> dict:
+    def model_dump(self, **kwargs: Any) -> dict:
         """Return dictionary representation of output parser."""
-        output_parser_dict = super().dict(**kwargs)
+        output_parser_dict = super().model_dump(**kwargs)
         with contextlib.suppress(NotImplementedError):
             output_parser_dict["_type"] = self._type
         return output_parser_dict
diff --git a/libs/core/langchain_core/output_parsers/json.py b/libs/core/langchain_core/output_parsers/json.py
@@ -9,6 +9,7 @@
 import jsonpatch  # type: ignore[import-untyped]
 import pydantic
 from pydantic import SkipValidation
+from pydantic.v1 import BaseModel
 
 from langchain_core.exceptions import OutputParserException
 from langchain_core.output_parsers.format_instructions import JSON_FORMAT_INSTRUCTIONS
@@ -19,16 +20,9 @@
     parse_json_markdown,
     parse_partial_json,
 )
-from langchain_core.utils.pydantic import IS_PYDANTIC_V1
 
-if IS_PYDANTIC_V1:
-    PydanticBaseModel = pydantic.BaseModel
-
-else:
-    from pydantic.v1 import BaseModel
-
-    # Union type needs to be last assignment to PydanticBaseModel to make mypy happy.
-    PydanticBaseModel = Union[BaseModel, pydantic.BaseModel]  # type: ignore[assignment,misc]
+# Union type needs to be last assignment to PydanticBaseModel to make mypy happy.
+PydanticBaseModel = Union[BaseModel, pydantic.BaseModel]
 
 TBaseModel = TypeVar("TBaseModel", bound=PydanticBaseModel)
 
diff --git a/libs/core/langchain_core/output_parsers/openai_functions.py b/libs/core/langchain_core/output_parsers/openai_functions.py
@@ -7,6 +7,7 @@
 
 import jsonpatch  # type: ignore[import-untyped]
 from pydantic import BaseModel, model_validator
+from pydantic.v1 import BaseModel as BaseModelV1
 from typing_extensions import override
 
 from langchain_core.exceptions import OutputParserException
@@ -274,10 +275,13 @@ def parse_result(self, result: list[Generation], *, partial: bool = False) -> An
                 pydantic_schema = self.pydantic_schema[fn_name]
             else:
                 pydantic_schema = self.pydantic_schema
-            if hasattr(pydantic_schema, "model_validate_json"):
+            if isinstance(pydantic_schema, BaseModel):
                 pydantic_args = pydantic_schema.model_validate_json(_args)
-            else:
+            elif isinstance(pydantic_schema, BaseModelV1):
                 pydantic_args = pydantic_schema.parse_raw(_args)
+            else:
+                msg = f"Unsupported pydantic schema: {pydantic_schema}"
+                raise ValueError(msg)
         return pydantic_args
 
 
diff --git a/libs/core/langchain_core/output_parsers/pydantic.py b/libs/core/langchain_core/output_parsers/pydantic.py
@@ -11,7 +11,6 @@
 from langchain_core.output_parsers import JsonOutputParser
 from langchain_core.outputs import Generation
 from langchain_core.utils.pydantic import (
-    IS_PYDANTIC_V2,
     PydanticBaseModel,
     TBaseModel,
 )
@@ -24,22 +23,16 @@ class PydanticOutputParser(JsonOutputParser, Generic[TBaseModel]):
     """The pydantic model to parse."""
 
     def _parse_obj(self, obj: dict) -> TBaseModel:
-        if IS_PYDANTIC_V2:
-            try:
-                if issubclass(self.pydantic_object, pydantic.BaseModel):
-                    return self.pydantic_object.model_validate(obj)
-                if issubclass(self.pydantic_object, pydantic.v1.BaseModel):
-                    return self.pydantic_object.parse_obj(obj)
-                msg = f"Unsupported model version for PydanticOutputParser: \
-                            {self.pydantic_object.__class__}"
-                raise OutputParserException(msg)
-            except (pydantic.ValidationError, pydantic.v1.ValidationError) as e:
-                raise self._parser_exception(e, obj) from e
-        else:  # pydantic v1
-            try:
+        try:
+            if issubclass(self.pydantic_object, pydantic.BaseModel):
+                return self.pydantic_object.model_validate(obj)
+            if issubclass(self.pydantic_object, pydantic.v1.BaseModel):
                 return self.pydantic_object.parse_obj(obj)
-            except pydantic.ValidationError as e:
-                raise self._parser_exception(e, obj) from e
+            msg = f"Unsupported model version for PydanticOutputParser: \
+                        {self.pydantic_object.__class__}"
+            raise OutputParserException(msg)
+        except (pydantic.ValidationError, pydantic.v1.ValidationError) as e:
+            raise self._parser_exception(e, obj) from e
 
     def _parser_exception(
         self, e: Exception, json_object: dict
diff --git a/libs/core/langchain_core/output_parsers/transform.py b/libs/core/langchain_core/output_parsers/transform.py
@@ -125,7 +125,7 @@ def _transform(self, input: Iterator[Union[str, BaseMessage]]) -> Iterator[Any]:
                 chunk_gen = ChatGenerationChunk(message=chunk)
             elif isinstance(chunk, BaseMessage):
                 chunk_gen = ChatGenerationChunk(
-                    message=BaseMessageChunk(**chunk.dict())
+                    message=BaseMessageChunk(**chunk.model_dump())
                 )
             else:
                 chunk_gen = GenerationChunk(text=chunk)
@@ -151,7 +151,7 @@ async def _atransform(
                 chunk_gen = ChatGenerationChunk(message=chunk)
             elif isinstance(chunk, BaseMessage):
                 chunk_gen = ChatGenerationChunk(
-                    message=BaseMessageChunk(**chunk.dict())
+                    message=BaseMessageChunk(**chunk.model_dump())
                 )
             else:
                 chunk_gen = GenerationChunk(text=chunk)
diff --git a/libs/core/langchain_core/prompts/base.py b/libs/core/langchain_core/prompts/base.py
@@ -331,7 +331,7 @@ def _prompt_type(self) -> str:
         """Return the prompt type key."""
         raise NotImplementedError
 
-    def dict(self, **kwargs: Any) -> dict:
+    def model_dump(self, **kwargs: Any) -> dict:
         """Return dictionary representation of prompt.
 
         Args:
@@ -369,7 +369,7 @@ def save(self, file_path: Union[Path, str]) -> None:
             raise ValueError(msg)
 
         # Fetch dictionary to save
-        prompt_dict = self.dict()
+        prompt_dict = self.model_dump()
         if "_type" not in prompt_dict:
             msg = f"Prompt {self} does not support saving."
             raise NotImplementedError(msg)
diff --git a/libs/core/langchain_core/pydantic_v1/__init__.py b/libs/core/langchain_core/pydantic_v1/__init__.py
@@ -18,7 +18,7 @@
 try:
     from pydantic.v1 import *  # noqa: F403
 except ImportError:
-    from pydantic import *  # type: ignore[assignment,no-redef] # noqa: F403
+    from pydantic import *  # type: ignore[assignment,deprecated,no-redef] # noqa: F403
 
 
 try:
diff --git a/libs/core/langchain_core/tools/base.py b/libs/core/langchain_core/tools/base.py
@@ -543,10 +543,13 @@ def _parse_input(
                     )
                     raise ValueError(msg)
                 key_ = next(iter(get_fields(input_args).keys()))
-                if hasattr(input_args, "model_validate"):
+                if issubclass(input_args, BaseModel):
                     input_args.model_validate({key_: tool_input})
-                else:
+                elif issubclass(input_args, BaseModelV1):
                     input_args.parse_obj({key_: tool_input})
+                else:
+                    msg = f"args_schema must be a Pydantic BaseModel, got {input_args}"
+                    raise TypeError(msg)
             return tool_input
         if input_args is not None:
             if isinstance(input_args, dict):
diff --git a/libs/core/langchain_core/tracers/schemas.py b/libs/core/langchain_core/tracers/schemas.py
@@ -2,8 +2,8 @@
 
 from __future__ import annotations
 
-import datetime
 import warnings
+from datetime import datetime, timezone
 from typing import Any, Optional
 from uuid import UUID
 
@@ -32,7 +32,7 @@ def RunTypeEnum() -> type[RunTypeEnumDep]:  # noqa: N802
 class TracerSessionV1Base(BaseModelV1):
     """Base class for TracerSessionV1."""
 
-    start_time: datetime.datetime = FieldV1(default_factory=datetime.datetime.utcnow)
+    start_time: datetime = FieldV1(default_factory=lambda: datetime.now(timezone.utc))
     name: Optional[str] = None
     extra: Optional[dict[str, Any]] = None
 
@@ -69,8 +69,8 @@ class BaseRun(BaseModelV1):
 
     uuid: str
     parent_uuid: Optional[str] = None
-    start_time: datetime.datetime = FieldV1(default_factory=datetime.datetime.utcnow)
-    end_time: datetime.datetime = FieldV1(default_factory=datetime.datetime.utcnow)
+    start_time: datetime = FieldV1(default_factory=lambda: datetime.now(timezone.utc))
+    end_time: datetime = FieldV1(default_factory=lambda: datetime.now(timezone.utc))
     extra: Optional[dict[str, Any]] = None
     execution_order: int
     child_execution_order: int
diff --git a/libs/core/langchain_core/utils/pydantic.py b/libs/core/langchain_core/utils/pydantic.py
diff --git a/libs/core/tests/unit_tests/output_parsers/test_openai_tools.py b/libs/core/tests/unit_tests/output_parsers/test_openai_tools.py
diff --git a/libs/core/tests/unit_tests/test_tools.py b/libs/core/tests/unit_tests/test_tools.py
diff --git a/libs/core/tests/unit_tests/utils/test_pydantic.py b/libs/core/tests/unit_tests/utils/test_pydantic.py
diff --git a/libs/core/tests/unit_tests/utils/test_utils.py b/libs/core/tests/unit_tests/utils/test_utils.py
