diff --git a/libs/langgraph/langgraph/graph/schema_utils.py b/libs/langgraph/langgraph/graph/schema_utils.py index 83b5a1b58..9bbde01bc 100644 --- a/libs/langgraph/langgraph/graph/schema_utils.py +++ b/libs/langgraph/langgraph/graph/schema_utils.py @@ -1,6 +1,7 @@ import functools import logging import weakref +from dataclasses import is_dataclass from inspect import isclass from typing import ( Annotated, @@ -13,8 +14,8 @@ from typing import ( get_type_hints, ) -from pydantic import BaseModel -from pydantic.v1 import BaseModel as BaseModelV1 +from pydantic import BaseModel, ConfigDict, TypeAdapter +from typing_extensions import is_typeddict __all__ = ["SchemaCoercionMapper"] @@ -45,7 +46,7 @@ class SchemaCoercionMapper: def __init__( self, - schema: type[Any], + schema: type[BaseModel], type_hints: Optional[dict[str, Any]] = None, *, max_depth: int = 12, @@ -63,30 +64,12 @@ class SchemaCoercionMapper: else get_type_hints(schema, localns={schema.__name__: schema}) ) - if issubclass(schema, BaseModelV1): - self._fields = { - n: self.type_hints.get(n, f.annotation) - for n, f in schema.__fields__.items() - } - self._construct = schema.construct - unhandled_attrs = ( - "__pre_root_validators__", - "__post_root_validators__", - "__validators__", - ) - if any(getattr(schema, c, None) for c in unhandled_attrs): - self.coerce: Callable[[Any, Any], Union[BaseModelV1, BaseModel]] = ( - lambda v, _: schema(**v) - ) - else: - self.coerce = self._coerce - - elif issubclass(schema, BaseModel): + if issubclass(schema, BaseModel): self._fields = { n: self.type_hints.get(n, f.annotation) for n, f in schema.model_fields.items() } - self._construct: Callable[..., Any] = schema.model_construct # type: ignore + self._construct: Callable[..., Any] = schema.model_construct unhandled_attrs = ("validators", "field_validators", "root_validators") if (decorators := getattr(schema, "__pydantic_decorators__", None)) and any( getattr(decorators, attr, None) for attr in unhandled_attrs @@ -94,9 +77,8 @@ class SchemaCoercionMapper: self.coerce = lambda v, _: schema.model_validate(v) else: self.coerce = self._coerce - else: - raise TypeError("Schema is neither a Pydantic v1 nor v2 model.") + raise TypeError("Schema must be a Pydantic V2 model.") self._field_coercers: Optional[dict[str, Callable[[Any, int], Any]]] = None @@ -138,14 +120,12 @@ class SchemaCoercionMapper: if isclass(field_type): # This is needed bcs. of issubclass issues on older versions of python - is_class_ = True try: - is_bm_v2 = issubclass(field_type, BaseModel) + is_bm_subclass = issubclass(field_type, BaseModel) except TypeError: # python < 3.11 issue. - is_class_ = False - is_bm_v2 = False - if is_bm_v2 or (is_class_ and issubclass(field_type, BaseModelV1)): + is_bm_subclass = False + if is_bm_subclass: mapper = SchemaCoercionMapper(field_type, max_depth=depth - 1) return lambda v, d: mapper.coerce(v, d) if isinstance(v, dict) else v @@ -265,67 +245,18 @@ _IDENTITY_TYPES: tuple[type[Any], ...] = ( type(None), ) -try: - # Pydantic v2. - from pydantic import TypeAdapter +@functools.lru_cache(maxsize=2048) +def _adapter_for(tp: Any) -> Callable[[Any], Any]: # noqa: D401 try: - import pydantic.v1.types as v1_types_ - from pydantic.v1 import parse_obj_as - - v1_types = tuple( - v for k, v in vars(v1_types_).items() if k in v1_types_.__all__ + config = ( + None + if (issubclass(tp, BaseModel) or is_dataclass(tp) or is_typeddict(tp)) + else ConfigDict(arbitrary_types_allowed=True) ) - except ImportError: - v1_types = () - - def parse_obj_as(tp: Any, v: Any) -> Any: # type: ignore - return v - - try: - from pydantic.v1 import parse_obj_as - from pydantic.v1.main import create_model - except ImportError: - create_model = None # type: ignore - - def _get_v1_parser(tp: Any) -> Any: - if create_model is not None: - try: - parser = create_model( - f"ParsingModel[{tp}]", - __root__=(tp, ...), - ) - return lambda v: parser(__root__=v).__root__ # type: ignore - except RuntimeError: - return lambda v: v - return lambda v: parse_obj_as(tp, v) - - @functools.lru_cache(maxsize=2048) - def _adapter_for(tp: Any) -> Callable[[Any], Any]: # noqa: D401 - if tp in v1_types: - return _get_v1_parser(tp) - try: - return TypeAdapter( - tp, config={"arbitrary_types_allowed": True} - ).validate_python - except TypeError: - # Delayed classes like ConstrainedList - return _get_v1_parser(tp) - -except ImportError: - # Pydantic V1 - from pydantic.v1.main import create_model - - @functools.lru_cache(maxsize=2048) - def _adapter_for(tp: Any) -> Callable[[Any], Any]: # noqa: D401 - try: - parser = create_model( - f"ParsingModel[{tp}]", - __root__=(tp, ...), - ) - return lambda v: parser(__root__=v).__root__ # type: ignore - except RuntimeError: - return lambda v: v + except TypeError: + config = None + return TypeAdapter(tp, config=config).validate_python def _get_adapter(tp: Any) -> Callable[[Any], Any]: diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 9653e2282..a81e8831e 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -23,7 +23,6 @@ from typing import ( from langchain_core.runnables import Runnable, RunnableConfig from pydantic import BaseModel -from pydantic.v1 import BaseModel as BaseModelV1 from typing_extensions import Self from langgraph._api.deprecation import LangGraphDeprecationWarning @@ -625,7 +624,7 @@ class StateGraph(Graph): self.input if len(self.channels) > 1 and isclass(self.input) - and issubclass(self.input, (BaseModel, BaseModelV1)) + and issubclass(self.input, BaseModel) else None ), nodes={}, @@ -1011,7 +1010,7 @@ def _pick_mapper( if isclass(schema): if issubclass(schema, dict): return None - if issubclass(schema, (BaseModel, BaseModelV1)): + if issubclass(schema, BaseModel): return SchemaCoercionMapper(schema, type_hints=type_hints) return partial(_coerce_state, schema) @@ -1194,7 +1193,7 @@ def _get_schema( channels: dict, name: str, ) -> type[BaseModel]: - if isclass(typ) and issubclass(typ, (BaseModel, BaseModelV1)): + if isclass(typ) and issubclass(typ, BaseModel): return typ else: keys = list(schemas[typ].keys()) diff --git a/libs/langgraph/langgraph/utils/fields.py b/libs/langgraph/langgraph/utils/fields.py index d94c58a64..e13c697d8 100644 --- a/libs/langgraph/langgraph/utils/fields.py +++ b/libs/langgraph/langgraph/utils/fields.py @@ -3,7 +3,6 @@ from collections.abc import Generator, Sequence from typing import Annotated, Any, Optional, Union, get_type_hints from pydantic import BaseModel -from pydantic.v1 import BaseModel as BaseModelV1 from typing_extensions import NotRequired, ReadOnly, Required, get_origin # NOTE: this is redefined here separately from langgraph.constants @@ -158,12 +157,7 @@ def get_enhanced_type_hints( def get_update_as_tuples(input: Any, keys: Sequence[str]) -> list[tuple[str, Any]]: """Get Pydantic state update as a list of (key, value) tuples.""" - # Pydantic v1 - if isinstance(input, BaseModelV1): - keep: Optional[set[str]] = input.__fields_set__ - defaults = {k: v.default for k, v in input.__fields__.items()} - # Pydantic v2 - elif isinstance(input, BaseModel): + if isinstance(input, BaseModel): keep = input.model_fields_set defaults = {k: v.default for k, v in input.model_fields.items()} else: diff --git a/libs/langgraph/langgraph/utils/pydantic.py b/libs/langgraph/langgraph/utils/pydantic.py index 66f434e5b..b44c942cb 100644 --- a/libs/langgraph/langgraph/utils/pydantic.py +++ b/libs/langgraph/langgraph/utils/pydantic.py @@ -1,11 +1,181 @@ import sys import typing +import warnings +from contextlib import nullcontext from dataclasses import is_dataclass -from typing import Any, Optional, Union +from functools import lru_cache +from typing import ( + Any, + Optional, + Union, + cast, + overload, +) -import typing_extensions -from pydantic import BaseModel -from pydantic.v1 import BaseModel as BaseModelV1 +from pydantic import ( + BaseModel, + ConfigDict, + Field, + RootModel, +) +from pydantic import ( + create_model as _create_model_base, +) +from pydantic.fields import FieldInfo +from pydantic.json_schema import ( + DEFAULT_REF_TEMPLATE, + GenerateJsonSchema, + JsonSchemaMode, +) +from typing_extensions import TypedDict + + +@overload +def get_fields(model: type[BaseModel]) -> dict[str, FieldInfo]: ... + + +@overload +def get_fields(model: BaseModel) -> dict[str, FieldInfo]: ... + + +def get_fields( + model: Union[type[BaseModel], BaseModel], +) -> dict[str, FieldInfo]: + """Get the field names of a Pydantic model.""" + if hasattr(model, "model_fields"): + return model.model_fields + + if hasattr(model, "__fields__"): + return model.__fields__ # type: ignore[return-value] + msg = f"Expected a Pydantic model. Got {type(model)}" + raise TypeError(msg) + + +_SchemaConfig = ConfigDict( + arbitrary_types_allowed=True, frozen=True, protected_namespaces=() +) + +NO_DEFAULT = object() + + +def _create_root_model( + name: str, + type_: Any, + module_name: Optional[str] = None, + default_: object = NO_DEFAULT, +) -> type[BaseModel]: + """Create a base class.""" + + def schema( + cls: type[BaseModel], + by_alias: bool = True, # noqa: FBT001,FBT002 + ref_template: str = DEFAULT_REF_TEMPLATE, + ) -> dict[str, Any]: + # Complains about schema not being defined in superclass + schema_ = super(cls, cls).schema( # type: ignore[misc] + by_alias=by_alias, ref_template=ref_template + ) + schema_["title"] = name + return schema_ + + def model_json_schema( + cls: type[BaseModel], + by_alias: bool = True, # noqa: FBT001,FBT002 + ref_template: str = DEFAULT_REF_TEMPLATE, + schema_generator: type[GenerateJsonSchema] = GenerateJsonSchema, + mode: JsonSchemaMode = "validation", + ) -> dict[str, Any]: + # Complains about model_json_schema not being defined in superclass + schema_ = super(cls, cls).model_json_schema( # type: ignore[misc] + by_alias=by_alias, + ref_template=ref_template, + schema_generator=schema_generator, + mode=mode, + ) + schema_["title"] = name + return schema_ + + base_class_attributes = { + "__annotations__": {"root": type_}, + "model_config": ConfigDict(arbitrary_types_allowed=True), + "schema": classmethod(schema), + "model_json_schema": classmethod(model_json_schema), + "__module__": module_name or "langchain_core.runnables.utils", + } + + if default_ is not NO_DEFAULT: + base_class_attributes["root"] = default_ + with warnings.catch_warnings(): + custom_root_type = type(name, (RootModel,), base_class_attributes) + return cast("type[BaseModel]", custom_root_type) + + +@lru_cache(maxsize=256) +def _create_root_model_cached( + model_name: str, + type_: Any, + *, + module_name: Optional[str] = None, + default_: object = NO_DEFAULT, +) -> type[BaseModel]: + return _create_root_model( + model_name, type_, default_=default_, module_name=module_name + ) + + +@lru_cache(maxsize=256) +def _create_model_cached( + model_name: str, + /, + **field_definitions: Any, +) -> type[BaseModel]: + return _create_model_base( + model_name, + __config__=_SchemaConfig, + **_remap_field_definitions(field_definitions), + ) + + +# Reserved names should capture all the `public` names / methods that are +# used by BaseModel internally. This will keep the reserved names up-to-date. +# For reference, the reserved names are: +# "construct", "copy", "dict", "from_orm", "json", "parse_file", "parse_obj", +# "parse_raw", "schema", "schema_json", "update_forward_refs", "validate", +# "model_computed_fields", "model_config", "model_construct", "model_copy", +# "model_dump", "model_dump_json", "model_extra", "model_fields", +# "model_fields_set", "model_json_schema", "model_parametrized_name", +# "model_post_init", "model_rebuild", "model_validate", "model_validate_json", +# "model_validate_strings" +_RESERVED_NAMES = {key for key in dir(BaseModel) if not key.startswith("_")} + + +def _remap_field_definitions(field_definitions: dict[str, Any]) -> dict[str, Any]: + """This remaps fields to avoid colliding with internal pydantic fields.""" + + remapped = {} + for key, value in field_definitions.items(): + if key.startswith("_") or key in _RESERVED_NAMES: + # Let's add a prefix to avoid colliding with internal pydantic fields + if isinstance(value, FieldInfo): + msg = ( + f"Remapping for fields starting with '_' or fields with a name " + f"matching a reserved name {_RESERVED_NAMES} is not supported if " + f" the field is a pydantic Field instance. Got {key}." + ) + raise NotImplementedError(msg) + type_, default_ = value + remapped[f"private_{key}"] = ( + type_, + Field( + default=default_, + alias=key, + serialization_alias=key, + title=key.lstrip("_").replace("_", " ").title(), + ), + ) + else: + remapped[key] = value + return remapped def create_model( @@ -13,32 +183,70 @@ def create_model( *, field_definitions: Optional[dict[str, Any]] = None, root: Optional[Any] = None, -) -> Union[BaseModel, BaseModelV1]: +) -> type[BaseModel]: """Create a pydantic model with the given field definitions. + Attention: + Please do not use outside of langchain packages. This API + is subject to change at any time. + Args: model_name: The name of the model. + module_name: The name of the module where the model is defined. + This is used by Pydantic to resolve any forward references. field_definitions: The field definitions for the model. root: Type for a root model (RootModel) + + Returns: + Type[BaseModel]: The created model. """ - try: - # for langchain-core >= 0.3.0 - from langchain_core.utils.pydantic import create_model_v2 + field_definitions = field_definitions or {} - return create_model_v2( - model_name, - field_definitions=field_definitions, - root=root, - ) - except ImportError: - # for langchain-core < 0.3.0 - from langchain_core.runnables.utils import create_model + if root: + if field_definitions: + msg = ( + "When specifying __root__ no other " + f"fields should be provided. Got {field_definitions}" + ) + raise NotImplementedError(msg) - v1_kwargs = {} - if root is not None: - v1_kwargs["__root__"] = root + if isinstance(root, tuple): + kwargs = {"type_": root[0], "default_": root[1]} + else: + kwargs = {"type_": root} - return create_model(model_name, **v1_kwargs, **(field_definitions or {})) + try: + named_root_model = _create_root_model_cached(model_name, **kwargs) + except TypeError: + # something in the arguments into _create_root_model_cached is not hashable + named_root_model = _create_root_model( + model_name, + **kwargs, + ) + return named_root_model + + # No root, just field definitions + names = set(field_definitions.keys()) + + capture_warnings = False + + for name in names: + # Also if any non-reserved name is used (e.g., model_id or model_name) + if name.startswith("model"): + capture_warnings = True + + with warnings.catch_warnings() if capture_warnings else nullcontext(): + if capture_warnings: + warnings.filterwarnings(action="ignore") + try: + return _create_model_cached(model_name, **field_definitions) + except TypeError: + # something in field definitions is not hashable + return _create_model_base( + model_name, + __config__=_SchemaConfig, + **_remap_field_definitions(field_definitions), + ) def is_supported_by_pydantic(type_: Any) -> bool: @@ -51,14 +259,12 @@ def is_supported_by_pydantic(type_: Any) -> bool: if is_dataclass(type_): return True - # Pydantic does not support mixing .v1 and root namespaces, so - # we only check for BaseModel (not pydantic.v1.BaseModel). if isinstance(type_, type) and issubclass(type_, BaseModel): return True if hasattr(type_, "__orig_bases__"): for base in type_.__orig_bases__: - if base is typing_extensions.TypedDict: + if base is TypedDict: return True elif base is typing.TypedDict: # noqa: TID251 # ignoring TID251 since it's OK to use typing.TypedDict in this case. diff --git a/libs/langgraph/poetry.lock b/libs/langgraph/poetry.lock index a49144223..fa4507768 100644 --- a/libs/langgraph/poetry.lock +++ b/libs/langgraph/poetry.lock @@ -26,7 +26,6 @@ description = "Reusable constraint types to use with typing.Annotated" optional = false python-versions = ">=3.8" groups = ["main", "dev"] -markers = "python_version < \"4.0\"" files = [ {file = "annotated_types-0.7.0-py3-none-any.whl", hash = "sha256:1f02e8b43a8fbbc3f3e0d4f0f4bfc8131bcb4eebe8849b8e5c773f3a1c582a53"}, {file = "annotated_types-0.7.0.tar.gz", hash = "sha256:aff07c09a53a08bc8cfccb9c85b05f1aa9a2a6f23728d790723543408344ce89"}, @@ -2270,7 +2269,6 @@ description = "Data validation using Python type hints" optional = false python-versions = ">=3.8" groups = ["main", "dev"] -markers = "python_version < \"4.0\"" files = [ {file = "pydantic-2.9.2-py3-none-any.whl", hash = "sha256:f048cec7b26778210e28a0459867920654d48e5e62db0958433636cde4254f12"}, {file = "pydantic-2.9.2.tar.gz", hash = "sha256:d155cef71265d1e9807ed1c32b4c8deec042a44a50a4188b25ac67ecd81a9c0f"}, @@ -2295,7 +2293,6 @@ description = "Core functionality for Pydantic validation and serialization" optional = false python-versions = ">=3.8" groups = ["main", "dev"] -markers = "python_version < \"4.0\"" files = [ {file = "pydantic_core-2.23.4-cp310-cp310-macosx_10_12_x86_64.whl", hash = "sha256:b10bd51f823d891193d4717448fab065733958bdb6a6b351967bd349d48d5c9b"}, {file = "pydantic_core-2.23.4-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:4fc714bdbfb534f94034efaa6eadd74e5b93c8fa6315565a222f7b6f42ca1166"}, @@ -3306,7 +3303,6 @@ files = [ {file = "typing_extensions-4.12.2-py3-none-any.whl", hash = "sha256:04e5ca0351e0f3f85c6853954072df659d0d13fac324d0072316b67d7794700d"}, {file = "typing_extensions-4.12.2.tar.gz", hash = "sha256:1a7ead55c7e559dd4dee8856e3a88b41225abfe1ce8df57b7c13915fe121ffb8"}, ] -markers = {main = "python_version < \"4.0\""} [[package]] name = "tzdata" @@ -3677,4 +3673,4 @@ type = ["pytest-mypy"] [metadata] lock-version = "2.1" python-versions = ">=3.9" -content-hash = "0e0c3fc2d5a5348c8df102497c7221a1f052f37a788b315eb82dfc7d63423e17" +content-hash = "770dcaa5816fffb667b5e3639c0ee5add24df0a99bb9a89a6d8fb2c4a79fb185" diff --git a/libs/langgraph/pyproject.toml b/libs/langgraph/pyproject.toml index 21da6a86f..b63778442 100644 --- a/libs/langgraph/pyproject.toml +++ b/libs/langgraph/pyproject.toml @@ -14,6 +14,7 @@ langgraph-checkpoint = "^2.0.10" langgraph-sdk = { version = ">=0.1.42", python = "<4.0" } langgraph-prebuilt = { version = ">=0.1.8", python = "<4.0" } xxhash = "^3.5.0" +pydantic = { version = ">=2.7.4"} [tool.poetry.group.dev.dependencies] pytest = "^8.3.2" diff --git a/libs/langgraph/tests/agents.py b/libs/langgraph/tests/agents.py index 6e5ce595a..d639724ab 100644 --- a/libs/langgraph/tests/agents.py +++ b/libs/langgraph/tests/agents.py @@ -6,22 +6,18 @@ from pydantic import BaseModel # define these objects to avoid importing langchain_core.agents # and therefore avoid relying on core Pydantic version class AgentAction(BaseModel): + """ + Represents a request to execute an action by an agent. + + The action consists of the name of the tool to execute and the input to pass + to the tool. The log is used to pass along extra information about the action. + """ + tool: str tool_input: Union[str, dict] log: str type: Literal["AgentAction"] = "AgentAction" - model_config = { - "json_schema_extra": { - "description": ( - """Represents a request to execute an action by an agent. - -The action consists of the name of the tool to execute and the input to pass -to the tool. The log is used to pass along extra information about the action.""" - ) - } - } - class AgentFinish(BaseModel): """Final return value of an ActionAgent. @@ -32,12 +28,3 @@ class AgentFinish(BaseModel): return_values: dict log: str type: Literal["AgentFinish"] = "AgentFinish" - model_config = { - "json_schema_extra": { - "description": ( - """Final return value of an ActionAgent. - -Agents return an AgentFinish when they have reached a stopping condition.""" - ) - } - } diff --git a/libs/langgraph/tests/test_messages_state.py b/libs/langgraph/tests/test_messages_state.py index d88e6226b..882b0cfed 100644 --- a/libs/langgraph/tests/test_messages_state.py +++ b/libs/langgraph/tests/test_messages_state.py @@ -12,13 +12,11 @@ from langchain_core.messages import ( ToolMessage, ) from pydantic import BaseModel -from pydantic.v1 import BaseModel as BaseModelV1 from typing_extensions import TypedDict from langgraph.graph import add_messages from langgraph.graph.message import REMOVE_ALL_MESSAGES, MessagesState from langgraph.graph.state import END, START, StateGraph -from tests.conftest import IS_LANGCHAIN_CORE_030_OR_GREATER from tests.messages import _AnyIdHumanMessage _, CORE_MINOR, CORE_PATCH = (int(v) for v in langchain_core.__version__.split(".")) @@ -175,19 +173,11 @@ def test_delete_all(): assert result == expected_result -MESSAGES_STATE_SCHEMAS = [MessagesState] -if IS_LANGCHAIN_CORE_030_OR_GREATER: +class MessagesStatePydantic(BaseModel): + messages: Annotated[list[AnyMessage], add_messages] - class MessagesStatePydantic(BaseModel): - messages: Annotated[list[AnyMessage], add_messages] - MESSAGES_STATE_SCHEMAS.append(MessagesStatePydantic) -else: - - class MessagesStatePydanticV1(BaseModelV1): - messages: Annotated[list[AnyMessage], add_messages] - - MESSAGES_STATE_SCHEMAS.append(MessagesStatePydanticV1) +MESSAGES_STATE_SCHEMAS = [MessagesState, MessagesStatePydantic] @pytest.mark.parametrize("state_schema", MESSAGES_STATE_SCHEMAS) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 75d408f29..97e6786a9 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -33,6 +33,7 @@ from langchain_core.runnables import ( ) from langchain_core.runnables.graph import Edge from langsmith import traceable +from pydantic import BaseModel, ConfigDict, Field, ValidationError from pytest_mock import MockerFixture from syrupy import SnapshotAssertion from typing_extensions import TypedDict @@ -2596,14 +2597,12 @@ def test_in_one_fan_out_state_graph_waiting_edge_via_branch( @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1( +def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2( snapshot: SnapshotAssertion, mocker: MockerFixture, request: pytest.FixtureRequest, checkpointer_name: str, ) -> None: - from pydantic.v1 import BaseModel, ValidationError - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") setup = mocker.Mock() teardown = mocker.Mock() @@ -2642,8 +2641,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1( yo: int class State(BaseModel): - class Config: - arbitrary_types_allowed = True + model_config = ConfigDict(arbitrary_types_allowed=True) query: str inner: Annotated[InnerObject, lambda x, y: y] @@ -2651,197 +2649,6 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1( docs: Annotated[list[str], sorted_add] client: Annotated[httpx.Client, Context(make_httpx_client)] - class Input(BaseModel): - query: str - inner: InnerObject - - class Output(BaseModel): - answer: str - docs: list[str] - - class StateUpdate(BaseModel): - query: Optional[str] = None - answer: Optional[str] = None - docs: Optional[list[str]] = None - - class UpdateDocs34(BaseModel): - docs: list[str] = ["doc3", "doc4"] - - def rewrite_query(data: State) -> State: - assert isinstance(data.inner, InnerObject) - return {"query": f"query: {data.query}"} - - def analyzer_one(data: State) -> State: - assert isinstance(data.inner, InnerObject) - return StateUpdate(query=f"analyzed: {data.query}") - - def retriever_one(data: State) -> State: - return {"docs": ["doc1", "doc2"]} - - def retriever_two(data: State) -> State: - time.sleep(0.1) - return UpdateDocs34() - - def qa(data: State) -> State: - return {"answer": ",".join(data.docs)} - - def decider(data: State) -> str: - assert isinstance(data, State) - return "retriever_two" - - workflow = StateGraph(State, input=Input, output=Output) - - workflow.add_node("rewrite_query", rewrite_query) - workflow.add_node("analyzer_one", analyzer_one) - workflow.add_node("retriever_one", retriever_one) - workflow.add_node("retriever_two", retriever_two) - workflow.add_node("qa", qa) - - workflow.set_entry_point("rewrite_query") - workflow.add_edge("rewrite_query", "analyzer_one") - workflow.add_edge("analyzer_one", "retriever_one") - workflow.add_conditional_edges( - "rewrite_query", decider, {"retriever_two": "retriever_two"} - ) - workflow.add_edge(["retriever_one", "retriever_two"], "qa") - workflow.set_finish_point("qa") - - app = workflow.compile() - - if checkpointer_name == "memory": - assert app.get_graph().draw_mermaid(with_styles=False) == snapshot - assert app.get_input_jsonschema() == snapshot - assert app.get_output_jsonschema() == snapshot - - with pytest.raises(ValidationError), assert_ctx_once(): - app.invoke({"query": {}}) - - with assert_ctx_once(): - assert app.invoke({"query": "what is weather in sf", "inner": {"yo": 1}}) == { - "docs": ["doc1", "doc2", "doc3", "doc4"], - "answer": "doc1,doc2,doc3,doc4", - } - - with assert_ctx_once(): - assert [ - *app.stream({"query": "what is weather in sf", "inner": {"yo": 1}}) - ] == [ - {"rewrite_query": {"query": "query: what is weather in sf"}}, - {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, - {"retriever_two": {"docs": ["doc3", "doc4"]}}, - {"retriever_one": {"docs": ["doc1", "doc2"]}}, - {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - ] - - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_after=["retriever_one"], - ) - config = {"configurable": {"thread_id": "1"}} - - with assert_ctx_once(): - assert [ - c - for c in app_w_interrupt.stream( - {"query": "what is weather in sf", "inner": {"yo": 1}}, config - ) - ] == [ - {"rewrite_query": {"query": "query: what is weather in sf"}}, - {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, - {"retriever_two": {"docs": ["doc3", "doc4"]}}, - {"retriever_one": {"docs": ["doc1", "doc2"]}}, - {"__interrupt__": ()}, - ] - - with assert_ctx_once(): - assert [c for c in app_w_interrupt.stream(None, config)] == [ - {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - ] - - with assert_ctx_once(): - assert app_w_interrupt.update_state( - config, {"docs": ["doc5"]}, as_node="rewrite_query" - ) == { - "configurable": { - "thread_id": "1", - "checkpoint_id": AnyStr(), - "checkpoint_ns": "", - } - } - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2( - snapshot: SnapshotAssertion, - mocker: MockerFixture, - request: pytest.FixtureRequest, - checkpointer_name: str, -) -> None: - from pydantic import BaseModel, ConfigDict, Field, ValidationError - from pydantic.v1 import BaseModel as BaseModelV1 - - IS_V1 = BaseModel is BaseModelV1 - - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - setup = mocker.Mock() - teardown = mocker.Mock() - - @contextmanager - def assert_ctx_once() -> Iterator[None]: - assert setup.call_count == 0 - assert teardown.call_count == 0 - try: - yield - finally: - assert setup.call_count == 1 - assert teardown.call_count == 1 - setup.reset_mock() - teardown.reset_mock() - - @contextmanager - def make_httpx_client() -> Iterator[httpx.Client]: - setup() - with httpx.Client() as client: - try: - yield client - finally: - teardown() - - def sorted_add( - x: list[str], y: Union[list[str], list[tuple[str, str]]] - ) -> list[str]: - if isinstance(y[0], tuple): - for rem, _ in y: - x.remove(rem) - y = [t[1] for t in y] - return sorted(operator.add(x, y)) - - class InnerObject(BaseModel): - yo: int - - if IS_V1: - - class State(BaseModel): - class Config: - arbitrary_types_allowed = True - - query: str - inner: Annotated[InnerObject, lambda x, y: y] - answer: Optional[str] = None - docs: Annotated[list[str], sorted_add] - client: Annotated[httpx.Client, Context(make_httpx_client)] - - else: - - class State(BaseModel): - model_config = ConfigDict(arbitrary_types_allowed=True) - - query: str - inner: Annotated[InnerObject, lambda x, y: y] - answer: Optional[str] = None - docs: Annotated[list[str], sorted_add] - client: Annotated[httpx.Client, Context(make_httpx_client)] - class StateUpdate(BaseModel): query: Optional[str] = None answer: Optional[str] = None @@ -2966,8 +2773,6 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_inp request: pytest.FixtureRequest, checkpointer_name: str, ) -> None: - from pydantic import BaseModel - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") def sorted_add( @@ -4415,7 +4220,6 @@ def test_remove_message_from_node(): def test_xray_lance(snapshot: SnapshotAssertion): from langchain_core.messages import AnyMessage, HumanMessage - from pydantic import BaseModel, Field class Analyst(BaseModel): affiliation: str = Field( @@ -5533,8 +5337,6 @@ def test_dict_mixed_return() -> None: def test_command_pydantic_dataclass() -> None: - from pydantic import BaseModel - class PydanticState(BaseModel): foo: str @@ -6351,9 +6153,7 @@ def test_double_interrupt_subgraph( @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_multi_resume( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: +def test_multi_resume(request: pytest.FixtureRequest, checkpointer_name: str) -> None: checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") class ChildState(TypedDict): @@ -6362,11 +6162,11 @@ def test_multi_resume( human_inputs: list[str] def get_human_input(state: ChildState): - human_input = interrupt(state['prompt']) + human_input = interrupt(state["prompt"]) return { - 'human_input': human_input, - 'human_inputs': [human_input], + "human_input": human_input, + "human_inputs": [human_input], } child_graph = ( @@ -6385,13 +6185,13 @@ def test_multi_resume( return [ Send( "child_graph", - {'prompt': prompt}, + {"prompt": prompt}, ) - for prompt in state['prompts'] + for prompt in state["prompts"] ] def cleanup(state: ParentState): - assert len(state['human_inputs']) == len(state["prompts"]) + assert len(state["human_inputs"]) == len(state["prompts"]) parent_graph = ( StateGraph(ParentState) @@ -6404,21 +6204,19 @@ def test_multi_resume( ) thread_config: RunnableConfig = { - 'configurable': { - 'thread_id': uuid.uuid4(), + "configurable": { + "thread_id": uuid.uuid4(), }, } - prompts = ['a', 'b', 'c', 'd', 'e'] + prompts = ["a", "b", "c", "d", "e"] events = parent_graph.invoke( - {'prompts': prompts}, - thread_config, - stream_mode='values' + {"prompts": prompts}, thread_config, stream_mode="values" ) - assert len(events['__interrupt__']) == len(prompts) - interrupt_values = {i.value for i in events['__interrupt__']} + assert len(events["__interrupt__"]) == len(prompts) + interrupt_values = {i.value for i in events["__interrupt__"]} assert interrupt_values == set(prompts) resume_map: dict[str, str] = { @@ -6428,11 +6226,8 @@ def test_multi_resume( result = parent_graph.invoke(Command(resume=resume_map), thread_config) assert result == { - 'prompts': prompts, - 'human_inputs': [ - f"human input for prompt {prompt}" - for prompt in prompts - ], + "prompts": prompts, + "human_inputs": [f"human input for prompt {prompt}" for prompt in prompts], } @@ -7169,8 +6964,6 @@ def test_node_destinations() -> None: def test_pydantic_none_state_update() -> None: - from pydantic import BaseModel - class State(BaseModel): foo: Optional[str] @@ -7182,8 +6975,6 @@ def test_pydantic_none_state_update() -> None: def test_pydantic_state_update_command() -> None: - from pydantic import BaseModel - class State(BaseModel): foo: Optional[str] @@ -7215,8 +7006,6 @@ def test_pydantic_state_update_command() -> None: def test_pydantic_state_mutation() -> None: - from pydantic import BaseModel, Field - class Inner(BaseModel): a: int = 0 @@ -7249,8 +7038,6 @@ def test_pydantic_state_mutation() -> None: def test_pydantic_state_mutation_command() -> None: - from pydantic import BaseModel, Field - class Inner(BaseModel): a: int = 0 @@ -7529,8 +7316,6 @@ def test_interrupt_subgraph_reenter_checkpointer_true( def test_empty_invoke() -> None: - from pydantic import BaseModel - def reducer_merge_dicts( dict1: dict[Any, Any], dict2: dict[Any, Any] ) -> dict[Any, Any]: @@ -7582,8 +7367,6 @@ def test_empty_invoke() -> None: def test_parallel_interrupts( request: pytest.FixtureRequest, checkpointer_name: str ) -> None: - from pydantic import BaseModel, Field - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") # --- CHILD GRAPH --- @@ -7759,8 +7542,6 @@ def test_parallel_interrupts( def test_parallel_interrupts_double( request: pytest.FixtureRequest, checkpointer_name: str ) -> None: - from pydantic import BaseModel, Field - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") # --- CHILD GRAPH --- @@ -8214,8 +7995,6 @@ def test_batch_update_as_input( def test_migration_graph(snapshot: SnapshotAssertion) -> None: - from pydantic import BaseModel - class DummyState(BaseModel): pass_count: int = 0 diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 55c04c853..661beec5b 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -26,6 +26,7 @@ import pytest from langchain_core.language_models import GenericFakeChatModel from langchain_core.runnables import RunnableConfig, RunnableLambda, RunnablePassthrough from langchain_core.utils.aiter import aclosing +from pydantic import BaseModel, ConfigDict, Field, ValidationError from pytest_mock import MockerFixture from syrupy import SnapshotAssertion from typing_extensions import TypedDict @@ -4663,16 +4664,9 @@ async def test_in_one_fan_out_state_graph_waiting_edge_via_branch( ] -@pytest.mark.parametrize("version", ["v1", "v2"]) -async def test_nested_pydantic_models(version: str) -> None: +async def test_nested_pydantic_models() -> None: """Test that nested Pydantic models are properly constructed from leaf nodes up.""" - # Define nested Pydantic models - if version == "v1": - from pydantic.v1 import BaseModel, Field - else: - from pydantic import BaseModel, Field - class NestedModel(BaseModel): value: int name: str @@ -4799,8 +4793,6 @@ async def test_nested_pydantic_models(version: str) -> None: async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class( snapshot: SnapshotAssertion, mocker: MockerFixture, checkpointer_name: str ) -> None: - from pydantic.v1 import BaseModel, ValidationError - setup = mocker.Mock() teardown = mocker.Mock() @@ -4835,8 +4827,7 @@ async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class( return sorted(operator.add(x, y)) class State(BaseModel): - class Config: - arbitrary_types_allowed = True + model_config = ConfigDict(arbitrary_types_allowed=True) query: str answer: Optional[str] = None @@ -4992,8 +4983,6 @@ async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class( async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2( snapshot: SnapshotAssertion, checkpointer_name: str ) -> None: - from pydantic import BaseModel, ValidationError - def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] ) -> list[str]: diff --git a/libs/langgraph/tests/test_pydantic.py b/libs/langgraph/tests/test_pydantic.py index 53e6fe4e7..8298970ba 100644 --- a/libs/langgraph/tests/test_pydantic.py +++ b/libs/langgraph/tests/test_pydantic.py @@ -8,7 +8,18 @@ import uuid from enum import Enum from typing import Annotated, Literal, Optional, Union -import pytest +from pydantic import ( + BaseModel, + ByteSize, + Field, + SecretStr, + confloat, + conint, + conlist, + constr, + field_validator, + model_validator, +) from langgraph.constants import END, START from langgraph.graph.state import StateGraph @@ -45,50 +56,10 @@ def test_is_supported_by_pydantic() -> None: assert is_supported_by_pydantic(PydanticModel) is True - if hasattr(pydantic, "v1"): - class PydanticModelV1(pydantic.v1.BaseModel): - x: int - - assert is_supported_by_pydantic(PydanticModelV1) is False - - assert is_supported_by_pydantic(int) is False - - -@pytest.mark.parametrize("version", ["v1", "v2"]) -def test_nested_pydantic_models(version: str) -> None: +def test_nested_pydantic_models() -> None: """Test that nested Pydantic models are properly constructed from leaf nodes up.""" - # Define nested Pydantic models - # Import necessary modules - - if version == "v1": - from pydantic.v1 import ( # type: ignore - BaseModel, - ByteSize, - Field, - SecretStr, - confloat, - conint, - conlist, - constr, - ) - else: - from pydantic import ( # type: ignore - BaseModel, - ByteSize, - Field, - SecretStr, - confloat, - conint, - conlist, - constr, - ) - from pydantic.v1 import BaseModel as BaseModelV1 - - if BaseModel is BaseModelV1: - pytest.skip("Cannot test pydantic v2 using installed version < 2") - class NestedModel(BaseModel): value: int name: str @@ -123,10 +94,7 @@ def test_nested_pydantic_models(version: str) -> None: name: str friends: list[str] = Field(default_factory=list) # IDs of friends - if version == "v2": - conlist_type = conlist(item_type=int, min_length=2, max_length=5) - else: - conlist_type = conlist(item_type=int, min_items=2, max_items=5) + conlist_type = conlist(item_type=int, min_length=2, max_length=5) class State(BaseModel): # Basic nested model tests @@ -314,8 +282,6 @@ def test_nested_pydantic_models(version: str) -> None: def test_pydantic_state_field_validator(): - from pydantic import BaseModel, field_validator, model_validator - class State(BaseModel): name: str text: str = "" @@ -346,32 +312,3 @@ def test_pydantic_state_field_validator(): g = builder.compile() res = g.invoke(input_state) assert res["text"] == "Hello, Validated John!" - - -def test_pydantic_v1_state_root_validator(): - from pydantic.v1 import BaseModel, root_validator - - class State(BaseModel): - name: str - text: str = "" - only_root: int = 13 - - @root_validator(pre=True) - @classmethod - def validate(cls, values: dict): - values["name"] = "Validated " + values["name"] - return values | {"only_root": 396} - - input_state = {"name": "John"} - - def process_node(state: State): - assert State(**input_state) == state - return {"text": "Hello, " + state.name + "!"} - - builder = StateGraph(state_schema=State) - builder.add_node("process", process_node) - builder.add_edge(START, "process") - builder.add_edge("process", END) - g = builder.compile() - res = g.invoke(input_state) - assert res["text"] == "Hello, Validated John!" diff --git a/libs/langgraph/tests/test_state.py b/libs/langgraph/tests/test_state.py index 1e8c2c623..011dab41e 100644 --- a/libs/langgraph/tests/test_state.py +++ b/libs/langgraph/tests/test_state.py @@ -7,7 +7,7 @@ from typing import Annotated as Annotated2 import pytest from langchain_core.runnables import RunnableConfig, RunnableLambda -from pydantic.v1 import BaseModel +from pydantic import BaseModel from typing_extensions import NotRequired, Required, TypedDict from langgraph.graph.state import StateGraph, _get_node_name, _warn_invalid_state_schema