Import TypedDict from typing_extensions, add ruff rule (#2910)

Within `libs/langgraph`, change all `TypedDict` imports to come from
`typing_extensions` rather than `typing`, as `pydantic` doesn't like the
latter.

Additionally, add a ruff rule to ban these imports too (so this doesn't
regress).

Solves  #2909.

libs/langgraph/bench/fanout_to_subgraph.py

@@ -1,5 +1,7 @@
 import operator
-from typing import Annotated, TypedDict
+from typing import Annotated
+
+from typing_extensions import TypedDict
 
 from langgraph.constants import END, START, Send
 from langgraph.graph.state import StateGraph

libs/langgraph/langgraph/graph/message.py

@@ -8,7 +8,6 @@ from typing import (
     Literal,
     Optional,
     Sequence,
-    TypedDict,
     Union,
     cast,
 )
@@ -22,6 +21,7 @@ from langchain_core.messages import (
     convert_to_messages,
     message_chunk_to_message,
 )
+from typing_extensions import TypedDict
 
 from langgraph.graph.state import StateGraph
 

libs/langgraph/langgraph/prebuilt/chat_agent_executor.py

@@ -381,7 +381,7 @@ def create_react_agent(
         Add complex prompt with custom graph state:
 
         ```pycon
-        >>> from typing import TypedDict
+        >>> from typing_extensions import TypedDict
         >>>
         >>> from langgraph.managed import IsLastStep
         >>> prompt = ChatPromptTemplate.from_messages(

libs/langgraph/langgraph/prebuilt/tool_node.py

@@ -601,7 +601,8 @@ def tools_condition(
         >>> from langgraph.prebuilt import ToolNode, tools_condition
         >>> from langgraph.graph.message import add_messages
         ...
-        >>> from typing import TypedDict, Annotated
+        >>> from typing import Annotated
+        >>> from typing_extensions import TypedDict
         ...
         >>> @tool
         >>> def divide(a: float, b: float) -> int:

libs/langgraph/langgraph/prebuilt/tool_validator.py

@@ -74,7 +74,8 @@ class ValidationNode(RunnableCallable):
 
     Examples:
         Example usage for re-prompting the model to generate a valid response:
-        >>> from typing import Literal, Annotated, TypedDict
+        >>> from typing import Literal, Annotated
+        >>> from typing_extensions import TypedDict
         ...
         >>> from langchain_anthropic import ChatAnthropic
         >>> from pydantic import BaseModel, validator

libs/langgraph/langgraph/pregel/debug.py

@@ -10,13 +10,13 @@ from typing import (
     Mapping,
     Optional,
     Sequence,
-    TypedDict,
     Union,
 )
 from uuid import UUID
 
 from langchain_core.runnables.config import RunnableConfig
 from langchain_core.utils.input import get_bolded_text, get_colored_text
+from typing_extensions import TypedDict
 
 from langgraph.channels.base import BaseChannel
 from langgraph.checkpoint.base import Checkpoint, CheckpointMetadata, PendingWrite

libs/langgraph/langgraph/types.py

@@ -13,14 +13,13 @@ from typing import (
     Optional,
     Sequence,
     Type,
-    TypedDict,
     TypeVar,
     Union,
     cast,
 )
 
 from langchain_core.runnables import Runnable, RunnableConfig
-from typing_extensions import Self
+from typing_extensions import Self, TypedDict
 
 from langgraph.checkpoint.base import (
     BaseCheckpointSaver,
@@ -373,7 +372,8 @@ def interrupt(value: Any) -> Any:
     Example:
         ```python
         import uuid
-        from typing import TypedDict, Optional
+        from typing import Optional
+        from typing_extensions import TypedDict
 
         from langgraph.checkpoint.memory import MemorySaver
         from langgraph.constants import START

libs/langgraph/pyproject.toml

@@ -38,7 +38,7 @@ py-spy = "^0.3.14"
 types-requests = "^2.32.0.20240914"
 
 [tool.ruff]
-lint.select = [ "E", "F", "I" ]
+lint.select = [ "E", "F", "I", "TID251" ]
 lint.ignore = [ "E501" ]
 line-length = 88
 indent-width = 4
@@ -52,6 +52,9 @@ line-ending = "auto"
 docstring-code-format = false
 docstring-code-line-length = "dynamic"
 
+[tool.ruff.lint.flake8-tidy-imports.banned-api]
+"typing.TypedDict".msg = "Use typing_extensions.TypedDict instead."
+
 [tool.mypy]
 # https://mypy.readthedocs.io/en/stable/config_file.html
 disallow_untyped_defs = "True"

libs/langgraph/tests/test_interruption.py

@@ -1,7 +1,6 @@
-from typing import TypedDict
-
 import pytest
 from pytest_mock import MockerFixture
+from typing_extensions import TypedDict
 
 from langgraph.graph import END, START, StateGraph
 from tests.conftest import (

libs/langgraph/tests/test_large_cases.py

@@ -4,13 +4,14 @@ import re
 import time
 from contextlib import contextmanager
 from dataclasses import replace
-from typing import Annotated, Any, Iterator, Literal, Optional, TypedDict, Union, cast
+from typing import Annotated, Any, Iterator, Literal, Optional, Union, cast
 
 import httpx
 import pytest
 from langchain_core.runnables import RunnableConfig, RunnableMap, RunnablePick
 from pytest_mock import MockerFixture
 from syrupy import SnapshotAssertion
+from typing_extensions import TypedDict
 
 from langgraph.channels.context import Context
 from langgraph.channels.last_value import LastValue

libs/langgraph/tests/test_large_cases_async.py

@@ -9,7 +9,6 @@ from typing import (
     AsyncIterator,
     Literal,
     Optional,
-    TypedDict,
     Union,
     cast,
 )
@@ -21,6 +20,7 @@ from langchain_core.runnables import RunnableConfig, RunnablePick
 from pydantic import BaseModel
 from pytest_mock import MockerFixture
 from syrupy import SnapshotAssertion
+from typing_extensions import TypedDict
 
 from langgraph.channels.context import Context
 from langgraph.channels.last_value import LastValue

libs/langgraph/tests/test_pregel.py

@@ -21,7 +21,6 @@ from typing import (
     Optional,
     Sequence,
     Tuple,
-    TypedDict,
     Union,
     get_type_hints,
 )
@@ -36,6 +35,7 @@ from langchain_core.runnables import (
 from langsmith import traceable
 from pytest_mock import MockerFixture
 from syrupy import SnapshotAssertion
+from typing_extensions import TypedDict
 
 from langgraph.channels.base import BaseChannel
 from langgraph.channels.binop import BinaryOperatorAggregate

libs/langgraph/tests/test_pregel_async.py

@@ -19,7 +19,6 @@ from typing import (
     Literal,
     Optional,
     Tuple,
-    TypedDict,
     Union,
 )
 from uuid import UUID
@@ -34,6 +33,7 @@ from langchain_core.runnables import (
 from langchain_core.utils.aiter import aclosing
 from pytest_mock import MockerFixture
 from syrupy import SnapshotAssertion
+from typing_extensions import TypedDict
 
 from langgraph.channels.base import BaseChannel
 from langgraph.channels.binop import BinaryOperatorAggregate

libs/langgraph/tests/test_tracing_interops.py

@@ -1,13 +1,14 @@
 import json
 import sys
 import time
-from typing import Any, Callable, Tuple, TypedDict, TypeVar
+from typing import Any, Callable, Tuple, TypeVar
 from unittest.mock import MagicMock
 
 import langsmith as ls
 import pytest
 from langchain_core.runnables import RunnableConfig
 from langchain_core.tracers import LangChainTracer
+from typing_extensions import TypedDict
 
 from langgraph.graph import StateGraph
 

libs/langgraph/tests/test_utils.py

@@ -9,7 +9,6 @@ from typing import (
     List,
     Literal,
     Optional,
-    TypedDict,
     TypeVar,
     Union,
 )
@@ -17,7 +16,7 @@ from unittest.mock import patch
 
 import langsmith
 import pytest
-from typing_extensions import Annotated, NotRequired, Required
+from typing_extensions import Annotated, NotRequired, Required, TypedDict
 
 from langgraph.graph import END, StateGraph
 from langgraph.graph.graph import CompiledGraph