0.3.10

- remove usage of require_at_least_one_of, we shouldn't be enforcing
presence of keys in inputs/updates, an empty dict is a valid
input/update
- ensure that values that were explcitly set/assigned in pydantic model
are saved even if equal to default value

.github/workflows/deploy_docs.yml

@@ -102,7 +102,14 @@ jobs:
       - name: Build llms-text
         run: make llms-text
       - name: Build site
-        run: make build-docs
+        run: |
+          # If this is main branch, then we want to download stats. we do this 
+          # with the env variable DOWNLOAD_STATS=true
+          if [ "${{ github.ref }}" == "refs/heads/main" ]; then
+            DOWNLOAD_STATS=true make build-docs
+          else
+            make build-docs
+          fi
         env:
           MKDOCS_GIT_COMMITTERS_APIKEY: ${{ secrets.MKDOCS_GIT_COMMITTERS_APIKEY }}
           OPENAI_API_KEY: sf-proj-1234567890 # fake placeholder, shouldn't actually be used

docs/Makefile

@@ -10,7 +10,15 @@ build-prebuilt:
 	# Use to create an update to date prebuilt page.
 	# Looks up download stats for each of the prebuilt packages and
 	# generates the final prebuilt page.
-	poetry run python -m _scripts.third_party_page.get_download_stats stats.yml
+	@if [ "$(DOWNLOAD_STATS)" = "true" ]; then \
+		set -x; \
+		poetry run python -m _scripts.third_party_page.get_download_stats stats.yml; \
+		set +x; \
+	else \
+		set -x; \
+		poetry run python -m _scripts.third_party_page.get_download_stats --fake stats.yml; \
+		set +x; \
+	fi
 	poetry run python -m _scripts.third_party_page.create_third_party_page stats.yml docs/prebuilt.md --language python
 
 build-docs: build-typedoc build-prebuilt

docs/_scripts/notebook_hooks.py

@@ -186,7 +186,7 @@ def _on_page_markdown_with_config(
 
     if remove_base64_images:
         # Remove base64 encoded images from markdown
-        markdown = re.sub(r"!\[.*?\]\(data:image/+;base64,[^\)]+\)", "", markdown)
+        markdown = re.sub(r"!\[.*?\]\(data:image/[^;]+;base64,[^)]+\)", "", markdown)
 
     return markdown
 

docs/_scripts/third_party_page/get_download_stats.py

@@ -30,10 +30,23 @@ PACKAGES_FILE = HERE / "packages.yml"
 PACKAGES = yaml.safe_load(PACKAGES_FILE.read_text())['packages']
 
 
-def _get_weekly_downloads(packages: list[Package]) -> list[ResolvedPackage]:
+def _get_weekly_downloads(packages: list[Package], fake: bool) -> list[ResolvedPackage]:
     """Retrieve the monthly download count for a list of packages from PyPIStats."""
     resolved_packages: list[ResolvedPackage] = []
 
+    if fake:
+        # To avoid making network requests during testing, return fake download counts
+        for package in packages:
+            resolved_packages.append(
+                {
+                    "name": package["name"],
+                    "repo": package["repo"],
+                    "weekly_downloads": -12345,
+                    "description": package["description"],
+                }
+            )
+        return resolved_packages
+
     for package in packages:
         # First check if package exists on PyPI
         pypi_url = f"https://pypi.org/pypi/{package['name']}/json"
@@ -88,13 +101,13 @@ def _get_weekly_downloads(packages: list[Package]) -> list[ResolvedPackage]:
 
 
 
-def main(output_file: str) -> None:
+def main(output_file: str, fake: bool) -> None:
     """Main function to generate package download information.
 
     Args:
         output_file: Path to the output YAML file.
     """
-    resolved_packages: list[ResolvedPackage] = _get_weekly_downloads(PACKAGES)
+    resolved_packages: list[ResolvedPackage] = _get_weekly_downloads(PACKAGES, fake)
 
     if not output_file.endswith(".yml"):
         raise ValueError("Output file must have a .yml extension")
@@ -115,6 +128,15 @@ if __name__ == "__main__":
             "downloads.yml"
         ),
     )
+    parser.add_argument(
+        "--fake",
+        default=False,
+        action="store_true",
+        help=(
+            "Generate fake download counts for testing purposes. "
+            "This option will not make any network requests."
+        ),
+    )
     args = parser.parse_args()
 
-    main(args.output_file)
+    main(args.output_file, args.fake)

docs/docs/cloud/how-tos/generative_ui_react.md

@@ -103,7 +103,8 @@ const AgentState = Annotation.Root({
 export const graph = new StateGraph(AgentState)
   .addNode("weather", async (state, config) => {
     // Provide the type of the component map to ensure
-    // type safety of `ui.push()` calls.
+    // type safety of `ui.push()` calls as well as
+    // pushing the messages to the `ui` and sending a custom event as well.
     const ui = typedUi<typeof ComponentMap>(config);
 
     const weather = await new ChatOpenAI({ model: "gpt-4o-mini" })
@@ -120,7 +121,7 @@ export const graph = new StateGraph(AgentState)
     // Emit UI elements with associated AI message
     ui.push({ name: "weather", props: weather }, { message: response });
 
-    return { messages: [response], ui: ui.items };
+    return { messages: [response] };
   })
   .addEdge("__start__", "weather")
   .compile();
@@ -217,7 +218,7 @@ By default `LoadExternalComponent` will use the `assistantId` from `useStream()`
 
 ### Access and interact with the thread state from the UI component
 
-You can access the thread state from the UI component by using the `useStreamContext` hook.
+You can access the thread state inside the UI component by using the `useStreamContext` hook.
 
 ```tsx
 import { useStreamContext } from "@langchain/langgraph-sdk/react-ui";
@@ -256,8 +257,14 @@ You can pass additional context to the client components by providing a `meta` p
 Then, you can access the `meta` prop in the UI component by using the `useStreamContext` hook.
 
 ```tsx
+import { useStreamContext } from "@langchain/langgraph-sdk/react-ui";
+
 const WeatherComponent = (props: { city: string }) => {
-  const { meta } = useStreamContext();
+  const { meta } = useStreamContext<
+    { city: string },
+    { MetaType: { userId?: string } }
+  >();
+
   return (
     <div>
       Weather for {props.city} (user: {meta?.userId})

docs/docs/how-tos/create-react-agent-memory.ipynb

@@ -122,20 +122,18 @@
     "\n",
     "# For this tutorial we will use custom tool that returns pre-defined values for weather in two cities (NYC & SF)\n",
     "\n",
-    "from typing import Literal\n",
-    "\n",
     "from langchain_core.tools import tool\n",
     "\n",
     "\n",
     "@tool\n",
-    "def get_weather(city: Literal[\"nyc\", \"sf\"]):\n",
+    "def get_weather(location: str) -> str:\n",
     "    \"\"\"Use this to get weather information.\"\"\"\n",
-    "    if city == \"nyc\":\n",
+    "    if any([city in location.lower() for city in [\"nyc\", \"new york city\"]]):\n",
     "        return \"It might be cloudy in nyc\"\n",
-    "    elif city == \"sf\":\n",
+    "    elif any([city in location.lower() for city in [\"sf\", \"san francisco\"]]):\n",
     "        return \"It's always sunny in sf\"\n",
     "    else:\n",
-    "        raise AssertionError(\"Unknown city\")\n",
+    "        return f\"I am not sure what the weather is in {location}\"\n",
     "\n",
     "\n",
     "tools = [get_weather]\n",

libs/langgraph/langgraph/graph/schema_utils.py

@@ -0,0 +1,162 @@
+import logging
+import weakref
+from inspect import isclass
+from typing import (
+    Any,
+    Callable,
+    Optional,
+    Type,
+    Union,
+    get_args,
+    get_origin,
+)
+
+from pydantic import BaseModel
+from pydantic.v1 import BaseModel as BaseModelV1
+from typing_extensions import Annotated
+
+logger = logging.getLogger(__name__)
+
+
+class SchemaCoercionMapper:
+    _cache: weakref.WeakKeyDictionary[Type[Any], dict[int, "SchemaCoercionMapper"]] = (
+        weakref.WeakKeyDictionary()
+    )
+
+    def __new__(cls, schema: Type[Any], max_depth: int = 5) -> "SchemaCoercionMapper":
+        if schema not in cls._cache:
+            cls._cache[schema] = {}
+        if max_depth in cls._cache[schema]:
+            return cls._cache[schema][max_depth]
+
+        inst = super().__new__(cls)
+        cls._cache[schema][max_depth] = inst
+        return inst
+
+    def __init__(self, schema: Type[Any], max_depth: int = 5):
+        if hasattr(self, "_inited"):
+            return
+        self._inited = True
+        self.schema = schema
+        self.max_depth = max_depth
+        if hasattr(schema, "model_fields") and hasattr(schema, "model_construct"):
+            self._fields = {n: f.annotation for n, f in schema.model_fields.items()}
+            self._construct = schema.model_construct
+        elif hasattr(schema, "__fields__") and callable(
+            getattr(schema, "construct", None)
+        ):
+            self._fields = {n: f.annotation for n, f in schema.__fields__.items()}
+            self._construct = schema.construct
+        else:
+            raise TypeError("Schema is neither valid Pydantic v1 nor v2 model.")
+        self._field_coercers: Optional[dict[str, Callable[[Any, Any], Any]]] = None
+
+    def __call__(self, input_data: Any, depth: Optional[int] = None) -> Any:
+        return self.coerce(input_data, depth)
+
+    def coerce(self, input_data: Any, depth: Optional[int] = None) -> Any:
+        if depth is None:
+            depth = self.max_depth
+        if not isinstance(input_data, dict) or depth <= 0:
+            return input_data
+        processed = {}
+        if self._field_coercers is None:
+            self._field_coercers = {
+                n: self._build_coercer(t) for n, t in self._fields.items()
+            }
+        for k, v in input_data.items():
+            fn = self._field_coercers.get(k)
+            processed[k] = fn(v, depth - 1) if fn else v
+        return self._construct(**processed)
+
+    def _build_coercer(self, field_type: Any) -> Callable[[Any, Any], Any]:
+        origin = get_origin(field_type)
+        if origin is Annotated:
+            real_type, *_ = get_args(field_type)
+            sub = self._build_coercer(real_type)
+            return lambda v, d: sub(v, d)
+        if isclass(field_type):
+            is_class_ = True
+            try:
+                is_base_model = issubclass(field_type, BaseModel)
+            except TypeError:
+                is_class_ = False
+                is_base_model = False
+
+            if is_base_model:
+                mapper = SchemaCoercionMapper(field_type, self.max_depth)
+                return lambda v, d: mapper.coerce(v, d) if isinstance(v, dict) else v
+            if is_class_ and issubclass(field_type, BaseModelV1):
+                mapper = SchemaCoercionMapper(field_type, self.max_depth)
+                return lambda v, d: mapper.coerce(v, d) if isinstance(v, dict) else v
+        if origin is list or field_type is list:
+            args = get_args(field_type)
+            if len(args) != 1:
+                return lambda v, d: v
+            sub = self._build_coercer(args[0])
+
+            def list_coercer(v: Any, d: Any) -> Any:
+                if not isinstance(v, (list, tuple)):
+                    raise TypeError(f"Expected list, got {type(v).__name__}")
+                return [sub(x, d - 1) for x in v]
+
+            return list_coercer
+        if origin is dict or field_type is dict:
+            args = get_args(field_type)
+            if len(args) != 2:
+
+                def plain_dict_coercer(v: Any, d: Any) -> Any:
+                    if not isinstance(v, dict):
+                        raise TypeError(f"Expected dict, got {type(v).__name__}")
+                    return v
+
+                return plain_dict_coercer
+            k_sub = self._build_coercer(args[0])
+            v_sub = self._build_coercer(args[1])
+
+            def dict_coercer(v: Any, d: Any) -> Any:
+                if not isinstance(v, dict):
+                    raise TypeError(f"Expected dict, got {type(v).__name__}")
+                return {k_sub(k, d - 1): v_sub(val, d - 1) for k, val in v.items()}
+
+            return dict_coercer
+
+        if origin is tuple:
+            targs = get_args(field_type)
+            if not targs:
+                return lambda v, d: v
+            subs = [self._build_coercer(a) for a in targs]
+
+            def tuple_coercer(v: Any, d: Any) -> Any:
+                if not isinstance(v, (list, tuple)):
+                    raise TypeError(f"Expected tuple-like, got {type(v).__name__}")
+                out = []
+                for i, sp in enumerate(subs):
+                    out.append(sp(v[i] if i < len(v) else None, d - 1))
+                return tuple(out)
+
+            return tuple_coercer
+        if origin is Union:
+            uargs = get_args(field_type)
+            subs, none_in_union = [], False
+            for arg in uargs:
+                if arg is type(None):
+                    none_in_union = True
+                else:
+                    subs.append(self._build_coercer(arg))
+
+            def union_coercer(v: Any, d: Any) -> Any:
+                if v is None and none_in_union:
+                    return None
+                err = None
+                for sp in subs:
+                    try:
+                        return sp(v, d - 1)
+                    except Exception as e:
+                        err = e
+                if err:
+                    raise err
+                return v
+
+            return union_coercer
+        return lambda v, d: v

libs/langgraph/langgraph/graph/state.py

@@ -50,6 +50,7 @@ from langgraph.graph.graph import (
     Graph,
     Send,
 )
+from langgraph.graph.schema_utils import SchemaCoercionMapper
 from langgraph.managed.base import (
     ChannelKeyPlaceholder,
     ChannelTypePlaceholder,
@@ -626,11 +627,13 @@ class StateGraph(Graph):
         compiled = CompiledStateGraph(
             builder=self,
             config_type=self.config_schema,
-            input_model=self.input
-            if len(self.channels) > 1
-            and isclass(self.input)
-            and issubclass(self.input, (BaseModel, BaseModelV1))
-            else None,
+            input_model=(
+                self.input
+                if len(self.channels) > 1
+                and isclass(self.input)
+                and issubclass(self.input, (BaseModel, BaseModelV1))
+                else None
+            ),
             nodes={},
             channels={
                 **self.channels,
@@ -759,23 +762,32 @@ class CompiledStateGraph(CompiledGraph):
                     else:
                         updates.extend(_get_updates(i) or ())
                 return updates
-            elif get_type_hints(type(input)):
+            elif (t := type(input)) and get_type_hints(t):
                 # Pydantic v2
-                if hasattr(input, "model_fields"):
+                if isinstance(input, BaseModel):
+                    keep: Optional[set[str]] = input.model_fields_set
                     defaults = {k: v.default for k, v in input.model_fields.items()}
                 # Pydantic v1
-                elif hasattr(input, "__fields__"):
-                    defaults = {k: v.default for k, v in input.__fields__.items()}
+                elif isinstance(input, BaseModelV1):
+                    keep = input.__fields_set__
+                    defaults = {k: v.default for k, v in t.__fields__.items()}
                 else:
+                    keep = None
                     defaults = {}
 
+                # NOTE: This behavior for Pydantic is somewhat inelegant,
+                # but we keep around for backwards compatibility
                 # if input is a Pydantic model, only update values
-                # that are different from the default values
+                # that are different from the default values or in the keep set
                 return [
                     (k, value)
                     for k in output_keys
                     if (value := getattr(input, k, MISSING)) is not MISSING
-                    and value != defaults.get(k)
+                    and (
+                        value is not None
+                        or defaults.get(k, MISSING) is not None
+                        or (keep is not None and k in keep)
+                    )
                 ]
             else:
                 msg = create_error_message(
@@ -801,7 +813,6 @@ class CompiledStateGraph(CompiledGraph):
                     ChannelWrite(
                         write_entries,
                         tags=[TAG_HIDDEN],
-                        require_at_least_one_of=output_keys,
                     ),
                 ],
             )
@@ -940,25 +951,14 @@ def _pick_mapper(
 ) -> Optional[Callable[[Any], Any]]:
     if state_keys == ["__root__"]:
         return None
-    if issubclass(schema, dict):
-        return None
-    if issubclass(schema, BaseModel):
-        return partial(_coerce_state_pydantic, schema)
-    if issubclass(schema, BaseModelV1):
-        return partial(_coerce_state_pydantic_v1, schema)
+    if isclass(schema):
+        if issubclass(schema, dict):
+            return None
+        if issubclass(schema, (BaseModel, BaseModelV1)):
+            return SchemaCoercionMapper(schema)
     return partial(_coerce_state, schema)
 
 
-def _coerce_state_pydantic(schema: Type[Any], input: dict[str, Any]) -> dict[str, Any]:
-    return schema.model_construct(**input)
-
-
-def _coerce_state_pydantic_v1(
-    schema: Type[Any], input: dict[str, Any]
-) -> dict[str, Any]:
-    return schema.construct(**input)
-
-
 def _coerce_state(schema: Type[Any], input: dict[str, Any]) -> dict[str, Any]:
     return schema(**input)
 

libs/langgraph/langgraph/pregel/debug.py

@@ -25,8 +25,10 @@ from langgraph.constants import (
     CONFIG_KEY_CHECKPOINT_NS,
     ERROR,
     INTERRUPT,
+    MISSING,
     NS_END,
     NS_SEP,
+    RETURN,
     TAG_HIDDEN,
 )
 from langgraph.pregel.io import read_channels
@@ -132,7 +134,9 @@ def map_debug_task_results(
             "id": task.id,
             "name": task.name,
             "error": next((w[1] for w in writes if w[0] == ERROR), None),
-            "result": [w for w in writes if w[0] in stream_channels_list],
+            "result": [
+                w for w in writes if w[0] in stream_channels_list or w[0] == RETURN
+            ],
             "interrupts": [asdict(w[1]) for w in writes if w[0] == INTERRUPT],
         },
     }
@@ -264,49 +268,63 @@ def tasks_w_writes(
 ) -> tuple[PregelTask, ...]:
     """Apply writes / subgraph states to tasks to be returned in a StateSnapshot."""
     pending_writes = pending_writes or []
-    return tuple(
-        PregelTask(
-            task.id,
-            task.name,
-            task.path,
-            next(
-                (
-                    exc
-                    for tid, n, exc in pending_writes
-                    if tid == task.id and n == ERROR
-                ),
-                None,
-            ),
-            tuple(
-                v for tid, n, v in pending_writes if tid == task.id and n == INTERRUPT
-            ),
-            states.get(task.id) if states else None,
+    out: list[PregelTask] = []
+    for task in tasks:
+        rtn = next(
             (
+                val
+                for tid, chan, val in pending_writes
+                if tid == task.id and chan == RETURN
+            ),
+            MISSING,
+        )
+        out.append(
+            PregelTask(
+                task.id,
+                task.name,
+                task.path,
                 next(
                     (
-                        val
-                        for tid, chan, val in pending_writes
-                        if tid == task.id and chan == output_keys
+                        exc
+                        for tid, n, exc in pending_writes
+                        if tid == task.id and n == ERROR
                     ),
                     None,
-                )
-                if isinstance(output_keys, str)
-                else {
-                    chan: val
-                    for tid, chan, val in pending_writes
-                    if tid == task.id
-                    and (
-                        chan == output_keys
-                        if isinstance(output_keys, str)
-                        else chan in output_keys
+                ),
+                tuple(
+                    v
+                    for tid, n, v in pending_writes
+                    if tid == task.id and n == INTERRUPT
+                ),
+                states.get(task.id) if states else None,
+                (
+                    rtn
+                    if rtn is not MISSING
+                    else next(
+                        (
+                            val
+                            for tid, chan, val in pending_writes
+                            if tid == task.id and chan == output_keys
+                        ),
+                        None,
                     )
-                }
+                    if isinstance(output_keys, str)
+                    else {
+                        chan: val
+                        for tid, chan, val in pending_writes
+                        if tid == task.id
+                        and (
+                            chan == output_keys
+                            if isinstance(output_keys, str)
+                            else chan in output_keys
+                        )
+                    }
+                )
+                if any(
+                    w[0] == task.id and w[1] not in (ERROR, INTERRUPT)
+                    for w in pending_writes
+                )
+                else None,
             )
-            if any(
-                w[0] == task.id and w[1] not in (ERROR, INTERRUPT)
-                for w in pending_writes
-            )
-            else None,
         )
-        for task in tasks
-    )
+    return tuple(out)

libs/langgraph/langgraph/pregel/read.py

@@ -201,7 +201,6 @@ class PregelNode(Runnable):
             writers[-2] = ChannelWrite(
                 writes=writers[-2].writes + writers[-1].writes,
                 tags=writers[-2].tags,
-                require_at_least_one_of=writers[-2].require_at_least_one_of,
             )
             writers.pop()
         return writers

libs/langgraph/langgraph/pregel/write.py

@@ -49,21 +49,18 @@ class ChannelWrite(RunnableCallable):
 
     writes: list[Union[ChannelWriteEntry, ChannelWriteTupleEntry, Send]]
     """Sequence of write entries or Send objects to write."""
-    require_at_least_one_of: Optional[Sequence[str]]
-    """If defined, at least one of these channels must be written to."""
 
     def __init__(
         self,
         writes: Sequence[Union[ChannelWriteEntry, ChannelWriteTupleEntry, Send]],
         *,
         tags: Optional[Sequence[str]] = None,
-        require_at_least_one_of: Optional[Sequence[str]] = None,
+        require_at_least_one_of: Optional[Sequence[str]] = None,  # ignored
     ):
         super().__init__(func=self._write, afunc=self._awrite, name=None, tags=tags)
         self.writes = cast(
             list[Union[ChannelWriteEntry, ChannelWriteTupleEntry, Send]], writes
         )
-        self.require_at_least_one_of = require_at_least_one_of
 
     def get_name(
         self, suffix: Optional[str] = None, *, name: Optional[str] = None
@@ -96,7 +93,6 @@ class ChannelWrite(RunnableCallable):
         self.do_write(
             config,
             writes,
-            self.require_at_least_one_of if input is not None else None,
         )
         return input
 
@@ -112,7 +108,6 @@ class ChannelWrite(RunnableCallable):
         self.do_write(
             config,
             writes,
-            self.require_at_least_one_of if input is not None else None,
         )
         return input
 
@@ -120,7 +115,7 @@ class ChannelWrite(RunnableCallable):
     def do_write(
         config: RunnableConfig,
         writes: Sequence[Union[ChannelWriteEntry, ChannelWriteTupleEntry, Send]],
-        require_at_least_one_of: Optional[Sequence[str]] = None,
+        require_at_least_one_of: Optional[Sequence[str]] = None,  # ignored
     ) -> None:
         # validate
         for w in writes:
@@ -151,12 +146,6 @@ class ChannelWrite(RunnableCallable):
                 tuples.append((w.channel, value))
             else:
                 raise ValueError(f"Invalid write entry: {w}")
-        # assert required channels
-        if require_at_least_one_of is not None:
-            if not {chan for chan, _ in tuples} & set(require_at_least_one_of):
-                raise InvalidUpdateError(
-                    f"Must write to at least one of {require_at_least_one_of}"
-                )
         write: TYPE_SEND = config[CONF][CONFIG_KEY_SEND]
         write(tuples)
 

libs/langgraph/langgraph/types.py

@@ -140,7 +140,7 @@ class PregelTask(NamedTuple):
     error: Optional[Exception] = None
     interrupts: tuple[Interrupt, ...] = ()
     state: Union[None, RunnableConfig, "StateSnapshot"] = None
-    result: Optional[dict[str, Any]] = None
+    result: Optional[Any] = None
 
 
 class PregelExecutableTask(NamedTuple):

libs/langgraph/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph"
-version = "0.3.8"
+version = "0.3.10"
 description = "Building stateful, multi-actor applications with LLMs"
 authors = []
 license = "MIT"

libs/langgraph/tests/test_pregel.py

@@ -2607,7 +2607,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1(
             arbitrary_types_allowed = True
 
         query: str
-        inner: InnerObject
+        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)]
@@ -2625,10 +2625,15 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1(
         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:
@@ -2636,7 +2641,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1(
 
     def retriever_two(data: State) -> State:
         time.sleep(0.1)
-        return {"docs": ["doc3", "doc4"]}
+        return UpdateDocs34()
 
     def qa(data: State) -> State:
         return {"answer": ",".join(data.docs)}
@@ -2732,7 +2737,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2(
     request: pytest.FixtureRequest,
     checkpointer_name: str,
 ) -> None:
-    from pydantic import BaseModel, ConfigDict, ValidationError
+    from pydantic import BaseModel, ConfigDict, Field, ValidationError
 
     checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
     setup = mocker.Mock()
@@ -2775,7 +2780,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2(
         model_config = ConfigDict(arbitrary_types_allowed=True)
 
         query: str
-        inner: InnerObject
+        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)]
@@ -2785,6 +2790,9 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2(
         answer: Optional[str] = None
         docs: Optional[list[str]] = None
 
+    class UpdateDocs34(BaseModel):
+        docs: list[str] = Field(default_factory=lambda: ["doc3", "doc4"])
+
     class Input(BaseModel):
         query: str
         inner: InnerObject
@@ -2794,9 +2802,11 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2(
         docs: list[str]
 
     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:
@@ -2804,7 +2814,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2(
 
     def retriever_two(data: State) -> State:
         time.sleep(0.1)
-        return {"docs": ["doc3", "doc4"]}
+        return UpdateDocs34()
 
     def qa(data: State) -> State:
         return {"answer": ",".join(data.docs)}
@@ -3027,6 +3037,123 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_inp
     }
 
 
+@pytest.mark.parametrize("version", ["v1", "v2"])
+def test_nested_pydantic_models(version: str) -> 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
+
+    # Forward reference model
+    class RecursiveModel(BaseModel):
+        value: str
+        child: Optional["RecursiveModel"] = None
+
+    # Discriminated union models
+    class Cat(BaseModel):
+        pet_type: Literal["cat"]
+        meow: str
+
+    class Dog(BaseModel):
+        pet_type: Literal["dog"]
+        bark: str
+
+    # Cyclic reference model
+    class Person(BaseModel):
+        id: str
+        name: str
+        friends: list[str] = Field(default_factory=list)  # IDs of friends
+
+    class State(BaseModel):
+        # Basic nested model tests
+        top_level: str
+        nested: NestedModel
+        optional_nested: Annotated[Optional[NestedModel], lambda x, y: y, "Foo"]
+        dict_nested: dict[str, NestedModel]
+        list_nested: Annotated[
+            Union[dict, list[dict[str, NestedModel]]], lambda x, y: (x or []) + [y]
+        ]
+        tuple_nested: tuple[str, NestedModel]
+        tuple_list_nested: list[tuple[int, NestedModel]]
+        complex_tuple: tuple[str, dict[str, tuple[int, NestedModel]]]
+
+        # Forward reference test
+        recursive: RecursiveModel
+
+        # Discriminated union test
+        pet: Union[Cat, Dog]
+
+        # Cyclic reference test
+        people: dict[str, Person]  # Map of ID -> Person
+
+    inputs = {
+        # Basic nested models
+        "top_level": "initial",
+        "nested": {"value": 42, "name": "test"},
+        "optional_nested": {"value": 10, "name": "optional"},
+        "dict_nested": {"a": {"value": 5, "name": "a"}},
+        "list_nested": [{"a": {"value": 6, "name": "b"}}],
+        "tuple_nested": ["tuple-key", {"value": 7, "name": "tuple-value"}],
+        "tuple_list_nested": [[1, {"value": 8, "name": "tuple-in-list"}]],
+        "complex_tuple": [
+            "complex",
+            {"nested": [9, {"value": 10, "name": "deep"}]},
+        ],
+        # Forward reference
+        "recursive": {"value": "parent", "child": {"value": "child", "child": None}},
+        # Discriminated union (using a cat in this case)
+        "pet": {"pet_type": "cat", "meow": "meow!"},
+        # Cyclic references
+        "people": {
+            "1": {
+                "id": "1",
+                "name": "Alice",
+                "friends": ["2", "3"],  # Alice is friends with Bob and Charlie
+            },
+            "2": {
+                "id": "2",
+                "name": "Bob",
+                "friends": ["1"],  # Bob is friends with Alice
+            },
+            "3": {
+                "id": "3",
+                "name": "Charlie",
+                "friends": ["1", "2"],  # Charlie is friends with Alice and Bob
+            },
+        },
+    }
+
+    update = {"top_level": "updated", "nested": {"value": 100, "name": "updated"}}
+
+    expected = State(**inputs)
+
+    def node_fn(state: State) -> dict:
+        assert state == expected
+        return update
+
+    builder = StateGraph(State)
+    builder.add_node("process", node_fn)
+    builder.set_entry_point("process")
+    builder.set_finish_point("process")
+    graph = builder.compile()
+
+    result = graph.invoke(inputs.copy())
+
+    assert result == {**inputs, **update}
+
+    new_inputs = inputs.copy()
+    new_inputs["list_nested"] = {"foo": "bar"}
+    expected = State(**new_inputs)
+    assert {**new_inputs, **update} == graph.invoke(new_inputs.copy())
+
+
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
 def test_in_one_fan_out_state_graph_waiting_edge_plus_regular(
     request: pytest.FixtureRequest, checkpointer_name: str
@@ -5550,37 +5677,6 @@ def test_command_goto_with_static_breakpoints(
     assert result == {"foo": "abc|node-1|node-2|node-2"}
 
 
-def test_nested_graph_state_error_handling():
-    """Test error handling when updating state in nested graphs."""
-
-    class State(TypedDict):
-        count: int
-
-    def child_node(state: State):
-        return {"count": state["count"] + 1}
-
-    child = StateGraph(State)
-    child.add_node("child", child_node)
-    child.add_edge(START, "child")
-
-    parent = StateGraph(State)
-    parent.add_node("child_graph", child.compile())
-    parent.add_edge(START, "child_graph")
-
-    app = parent.compile(checkpointer=MemorySaver())
-
-    # Test invalid state update on parent
-    with pytest.raises(InvalidUpdateError):
-        app.update_state({"configurable": {"thread_id": "1"}}, {"invalid_key": "value"})
-
-    # Test invalid state update on child
-    with pytest.raises(InvalidUpdateError):
-        app.update_state(
-            {"configurable": {"thread_id": "1", "checkpoint_ns": "child_graph"}},
-            {"invalid_key": "value"},
-        )
-
-
 def test_parallel_node_execution():
     """Test that parallel nodes execute concurrently."""
 
@@ -5821,8 +5917,267 @@ def test_falsy_return_from_task(
         interrupt("test")
 
     configurable = {"configurable": {"thread_id": str(uuid.uuid4())}}
-    graph.invoke({"a": 5}, configurable)
-    graph.invoke(Command(resume="123"), configurable)
+    assert [
+        chunk for chunk in graph.stream({"a": 5}, configurable, stream_mode="debug")
+    ] == [
+        {
+            "payload": {
+                "config": {
+                    "callbacks": None,
+                    "configurable": {
+                        "checkpoint_id": AnyStr(),
+                        "checkpoint_ns": "",
+                        "thread_id": AnyStr(),
+                    },
+                    "metadata": configurable["configurable"],
+                    "recursion_limit": 25,
+                    "tags": [],
+                },
+                "metadata": {
+                    "parents": {},
+                    "source": "input",
+                    "step": -1,
+                    "thread_id": AnyStr(),
+                    "writes": {
+                        "__start__": {
+                            "a": 5,
+                        },
+                    },
+                },
+                "next": [
+                    "graph",
+                ],
+                "parent_config": None,
+                "tasks": [
+                    {
+                        "id": AnyStr(),
+                        "interrupts": (),
+                        "name": "graph",
+                        "state": None,
+                    },
+                ],
+                "values": None,
+            },
+            "step": -1,
+            "timestamp": AnyStr(),
+            "type": "checkpoint",
+        },
+        {
+            "payload": {
+                "id": AnyStr(),
+                "input": {
+                    "a": 5,
+                },
+                "name": "graph",
+                "triggers": [
+                    "__start__",
+                ],
+            },
+            "step": 0,
+            "timestamp": AnyStr(),
+            "type": "task",
+        },
+        {
+            "payload": {
+                "id": AnyStr(),
+                "input": (
+                    (),
+                    {},
+                ),
+                "name": "falsy_task",
+                "triggers": [
+                    "__pregel_push",
+                ],
+            },
+            "step": 0,
+            "timestamp": AnyStr(),
+            "type": "task",
+        },
+        {
+            "payload": {
+                "error": None,
+                "id": AnyStr(),
+                "interrupts": [],
+                "name": "falsy_task",
+                "result": [
+                    (
+                        "__return__",
+                        False,
+                    ),
+                ],
+            },
+            "step": 0,
+            "timestamp": AnyStr(),
+            "type": "task_result",
+        },
+        {
+            "payload": {
+                "error": None,
+                "id": AnyStr(),
+                "interrupts": [
+                    {
+                        "ns": [
+                            AnyStr(),
+                        ],
+                        "resumable": True,
+                        "value": "test",
+                        "when": "during",
+                    },
+                ],
+                "name": "graph",
+                "result": [],
+            },
+            "step": 0,
+            "timestamp": AnyStr(),
+            "type": "task_result",
+        },
+    ]
+    assert [
+        c
+        for c in graph.stream(Command(resume="123"), configurable, stream_mode="debug")
+    ] == [
+        {
+            "payload": {
+                "config": {
+                    "callbacks": None,
+                    "configurable": {
+                        "checkpoint_id": AnyStr(),
+                        "checkpoint_ns": "",
+                        "thread_id": AnyStr(),
+                    },
+                    "metadata": configurable["configurable"],
+                    "recursion_limit": 25,
+                    "tags": [],
+                },
+                "metadata": {
+                    "parents": {},
+                    "source": "input",
+                    "step": -1,
+                    "thread_id": AnyStr(),
+                    "writes": {
+                        "__start__": {
+                            "a": 5,
+                        },
+                    },
+                },
+                "next": [
+                    "graph",
+                ],
+                "parent_config": None,
+                "tasks": [
+                    {
+                        "id": AnyStr(),
+                        "interrupts": (
+                            {
+                                "ns": [
+                                    AnyStr(),
+                                ],
+                                "resumable": True,
+                                "value": "test",
+                                "when": "during",
+                            },
+                        ),
+                        "name": "graph",
+                        "state": None,
+                    },
+                ],
+                "values": None,
+            },
+            "step": -1,
+            "timestamp": AnyStr(),
+            "type": "checkpoint",
+        },
+        {
+            "payload": {
+                "id": AnyStr(),
+                "input": {
+                    "a": 5,
+                },
+                "name": "graph",
+                "triggers": [
+                    "__start__",
+                ],
+            },
+            "step": 0,
+            "timestamp": AnyStr(),
+            "type": "task",
+        },
+        {
+            "payload": {
+                "id": AnyStr(),
+                "input": (
+                    (),
+                    {},
+                ),
+                "name": "falsy_task",
+                "triggers": [
+                    "__pregel_push",
+                ],
+            },
+            "step": 0,
+            "timestamp": AnyStr(),
+            "type": "task",
+        },
+        {
+            "payload": {
+                "error": None,
+                "id": AnyStr(),
+                "interrupts": [],
+                "name": "graph",
+                "result": [
+                    (
+                        "__end__",
+                        None,
+                    ),
+                ],
+            },
+            "step": 0,
+            "timestamp": AnyStr(),
+            "type": "task_result",
+        },
+        {
+            "payload": {
+                "config": {
+                    "callbacks": None,
+                    "configurable": {
+                        "checkpoint_id": AnyStr(),
+                        "checkpoint_ns": "",
+                        "thread_id": AnyStr(),
+                    },
+                    "metadata": configurable["configurable"],
+                    "recursion_limit": 25,
+                    "tags": [],
+                },
+                "metadata": {
+                    "parents": {},
+                    "source": "loop",
+                    "step": 0,
+                    "thread_id": AnyStr(),
+                    "writes": {
+                        "falsy_task": False,
+                        "graph": None,
+                    },
+                },
+                "next": [],
+                "parent_config": {
+                    "callbacks": None,
+                    "configurable": {
+                        "checkpoint_id": AnyStr(),
+                        "checkpoint_ns": "",
+                        "thread_id": AnyStr(),
+                    },
+                    "metadata": configurable["configurable"],
+                    "recursion_limit": 25,
+                    "tags": [],
+                },
+                "tasks": [],
+                "values": None,
+            },
+            "step": 0,
+            "timestamp": AnyStr(),
+            "type": "checkpoint",
+        },
+    ]
 
 
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
@@ -6912,3 +7267,53 @@ def test_interrupt_subgraph_reenter_checkpointer_true(
     }
     # confirm that we preserve the state values from the previous invocation
     assert bar_values == [None, "barbaz", "quxbaz"]
+
+
+def test_empty_invoke() -> None:
+    from pydantic import BaseModel
+
+    def reducer_merge_dicts(
+        dict1: dict[Any, Any], dict2: dict[Any, Any]
+    ) -> dict[Any, Any]:
+        merged = {**dict1, **dict2}
+        return merged
+
+    class SimpleGraphState(BaseModel):
+        x1: Annotated[list[str], operator.add] = []
+        x2: Annotated[dict[str, Any], reducer_merge_dicts] = {}
+
+    def update_x1_1(state: SimpleGraphState):
+        print(state)
+        return {"x1": ["111"]}
+
+    def update_x1_2(state: SimpleGraphState):
+        print(state)
+        state.x1.append("222")
+        return {"x1": ["222"]}
+
+    def update_x2_1(state: SimpleGraphState):
+        print(state)
+        return {"x2": {"111": 111}}
+
+    def update_x2_2(state: SimpleGraphState):
+        print(state)
+        return {"x2": {"222": 222}}
+
+    graph = StateGraph(SimpleGraphState)
+    graph.add_node("x1_1_node", update_x1_1)
+    graph.add_node("x1_2_node", update_x1_2)
+    graph.add_node("x2_1_node", update_x2_1)
+    graph.add_node("x2_2_node", update_x2_2)
+    graph.add_edge("x1_1_node", "x1_2_node")
+    graph.add_edge("x1_2_node", "x2_1_node")
+    graph.add_edge("x2_1_node", "x2_2_node")
+
+    graph.add_edge(START, "x1_1_node")
+    graph.add_edge("x2_2_node", END)
+
+    compiled = graph.compile()
+
+    assert compiled.invoke(SimpleGraphState()).get("x2") == {
+        "111": 111,
+        "222": 222,
+    }

libs/langgraph/tests/test_pregel_async.py

@@ -4511,6 +4511,116 @@ 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:
+    """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
+
+    # Forward reference model
+    class RecursiveModel(BaseModel):
+        value: str
+        child: Optional["RecursiveModel"] = None
+
+    # Discriminated union models
+    class Cat(BaseModel):
+        pet_type: Literal["cat"]
+        meow: str
+
+    class Dog(BaseModel):
+        pet_type: Literal["dog"]
+        bark: str
+
+    # Cyclic reference model
+    class Person(BaseModel):
+        id: str
+        name: str
+        friends: list[str] = Field(default_factory=list)  # IDs of friends
+
+    class State(BaseModel):
+        # Basic nested model tests
+        top_level: str
+        nested: NestedModel
+        optional_nested: Optional[NestedModel] = None
+        dict_nested: dict[str, NestedModel]
+        list_nested: Annotated[
+            Union[dict, list[dict[str, NestedModel]]], lambda x, y: (x or []) + [y]
+        ]
+        tuple_nested: tuple[str, NestedModel]
+        tuple_list_nested: list[tuple[int, NestedModel]]
+        complex_tuple: tuple[str, dict[str, tuple[int, NestedModel]]]
+
+        # Forward reference test
+        recursive: RecursiveModel
+
+        # Discriminated union test
+        pet: Union[Cat, Dog]
+
+        # Cyclic reference test
+        people: dict[str, Person]  # Map of ID -> Person
+
+    inputs = {
+        # Basic nested models
+        "top_level": "initial",
+        "nested": {"value": 42, "name": "test"},
+        "optional_nested": {"value": 10, "name": "optional"},
+        "dict_nested": {"a": {"value": 5, "name": "a"}},
+        "list_nested": [{"a": {"value": 6, "name": "b"}}],
+        "tuple_nested": ["tuple-key", {"value": 7, "name": "tuple-value"}],
+        "tuple_list_nested": [[1, {"value": 8, "name": "tuple-in-list"}]],
+        "complex_tuple": [
+            "complex",
+            {"nested": [9, {"value": 10, "name": "deep"}]},
+        ],
+        # Forward reference
+        "recursive": {"value": "parent", "child": {"value": "child", "child": None}},
+        # Discriminated union (using a cat in this case)
+        "pet": {"pet_type": "cat", "meow": "meow!"},
+        # Cyclic references
+        "people": {
+            "1": {
+                "id": "1",
+                "name": "Alice",
+                "friends": ["2", "3"],  # Alice is friends with Bob and Charlie
+            },
+            "2": {
+                "id": "2",
+                "name": "Bob",
+                "friends": ["1"],  # Bob is friends with Alice
+            },
+            "3": {
+                "id": "3",
+                "name": "Charlie",
+                "friends": ["1", "2"],  # Charlie is friends with Alice and Bob
+            },
+        },
+    }
+
+    update = {"top_level": "updated", "nested": {"value": 100, "name": "updated"}}
+
+    async def node_fn(state: State) -> dict:
+        assert state == State(**inputs)
+        return update
+
+    builder = StateGraph(State)
+    builder.add_node("process", node_fn)
+    builder.set_entry_point("process")
+    builder.set_finish_point("process")
+    graph = builder.compile()
+
+    result = await graph.ainvoke(inputs.copy())
+
+    assert result == {**inputs, **update}
+
+
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
 async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class(
     snapshot: SnapshotAssertion, mocker: MockerFixture, checkpointer_name: str
@@ -6544,39 +6654,6 @@ async def test_command_goto_with_static_breakpoints(checkpointer_name: str) -> N
         assert result == {"foo": "abc|node-1|node-2|node-2"}
 
 
-async def test_nested_graph_state_error_handling():
-    """Test error handling when updating state in nested graphs."""
-
-    class State(TypedDict):
-        count: int
-
-    def child_node(state: State):
-        return {"count": state["count"] + 1}
-
-    child = StateGraph(State)
-    child.add_node("child", child_node)
-    child.add_edge(START, "child")
-
-    parent = StateGraph(State)
-    parent.add_node("child_graph", child.compile())
-    parent.add_edge(START, "child_graph")
-
-    app = parent.compile(checkpointer=MemorySaver())
-
-    # Test invalid state update on parent
-    with pytest.raises(InvalidUpdateError):
-        await app.aupdate_state(
-            {"configurable": {"thread_id": "1"}}, {"invalid_key": "value"}
-        )
-
-    # Test invalid state update on child
-    with pytest.raises(InvalidUpdateError):
-        await app.aupdate_state(
-            {"configurable": {"thread_id": "1", "checkpoint_ns": "child_graph"}},
-            {"invalid_key": "value"},
-        )
-
-
 async def test_parallel_node_execution():
     """Test that parallel nodes execute concurrently."""
 

libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py

@@ -382,12 +382,11 @@ def create_react_agent(
         Use with a simple tool:
 
         ```pycon
-        >>> from datetime import datetime
         >>> from langchain_openai import ChatOpenAI
         >>> from langgraph.prebuilt import create_react_agent
 
 
-        ... def check_weather(location: str, at_time: datetime | None = None) -> str:
+        ... def check_weather(location: str) -> str:
         ...     '''Return the weather forecast for the specified location.'''
         ...     return f"It's always sunny in {location}"
         >>>
@@ -595,7 +594,7 @@ def create_react_agent(
 
         ```pycon
         >>> import time
-        ... def check_weather(location: str, at_time: datetime | None = None) -> float:
+        ... def check_weather(location: str) -> str:
         ...     '''Return the weather forecast for the specified location.'''
         ...     time.sleep(2)
         ...     return f"It's always sunny in {location}"

libs/sdk-js/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@langchain/langgraph-sdk",
-  "version": "0.0.54",
+  "version": "0.0.57",
   "description": "Client library for interacting with the LangGraph API",
   "type": "module",
   "packageManager": "yarn@1.22.19",

libs/sdk-js/src/client.ts

@@ -1022,13 +1022,23 @@ export class RunsClient<
    *
    * @param threadId The ID of the thread.
    * @param runId The ID of the run.
+   * @param options Additional options for controlling the stream behavior:
+   *   - signal: An AbortSignal that can be used to cancel the stream request
+   *   - cancelOnDisconnect: When true, automatically cancels the run if the client disconnects from the stream
+   *   - streamMode: Controls what types of events to receive from the stream (can be a single mode or array of modes)
+   *        Must be a subset of the stream modes passed when creating the run. Background runs default to having the union of all
+   *        stream modes enabled.
    * @returns An async generator yielding stream parts.
    */
   async *joinStream(
     threadId: string,
     runId: string,
     options?:
-      | { signal?: AbortSignal; cancelOnDisconnect?: boolean }
+      | {
+          signal?: AbortSignal;
+          cancelOnDisconnect?: boolean;
+          streamMode?: StreamMode | StreamMode[];
+        }
       | AbortSignal,
   ): AsyncGenerator<{ event: StreamEvent; data: any }> {
     const opts =
@@ -1043,7 +1053,10 @@ export class RunsClient<
         method: "GET",
         timeoutMs: null,
         signal: opts?.signal,
-        params: { cancel_on_disconnect: opts?.cancelOnDisconnect ? "1" : "0" },
+        params: {
+          cancel_on_disconnect: opts?.cancelOnDisconnect ? "1" : "0",
+          stream_mode: opts?.streamMode,
+        },
       }),
     );
 

libs/sdk-js/src/react-ui/server/server.ts

@@ -6,15 +6,34 @@ interface MessageLike {
   id?: string;
 }
 
-export const typedUi = <Decl extends Record<string, ElementType>>(config: {
-  writer?: (chunk: unknown) => void;
-  runId?: string;
-  metadata?: Record<string, unknown>;
-  tags?: string[];
-  runName?: string;
-}) => {
+/**
+ * Helper to send and persist UI messages. Accepts a map of component names to React components
+ * as type argument to provide type safety. Will also write to the `options?.stateKey` state.
+ *
+ * @param config LangGraphRunnableConfig
+ * @param options
+ * @returns
+ */
+export const typedUi = <Decl extends Record<string, ElementType>>(
+  config: {
+    writer?: (chunk: unknown) => void;
+    runId?: string;
+    metadata?: Record<string, unknown>;
+    tags?: string[];
+    runName?: string;
+    configurable?: {
+      __pregel_send?: (writes_: [string, unknown][]) => void;
+      [key: string]: unknown;
+    };
+  },
+  options?: {
+    /** The key to write the UI messages to. Defaults to `ui`. */
+    stateKey?: string;
+  },
+) => {
   type PropMap = { [K in keyof Decl]: ComponentPropsWithoutRef<Decl[K]> };
   let items: (UIMessage | RemoveUIMessage)[] = [];
+  const stateKey = options?.stateKey ?? "ui";
 
   const runId = (config.metadata?.run_id as string | undefined) ?? config.runId;
   if (!runId) throw new Error("run_id is required");
@@ -48,6 +67,7 @@ export const typedUi = <Decl extends Record<string, ElementType>>(config: {
     };
     items.push(evt);
     config.writer?.(evt);
+    config.configurable?.__pregel_send?.([[stateKey, evt]]);
     return evt;
   };
 
@@ -55,6 +75,7 @@ export const typedUi = <Decl extends Record<string, ElementType>>(config: {
     const evt: RemoveUIMessage = { type: "remove-ui", id };
     items.push(evt);
     config.writer?.(evt);
+    config.configurable?.__pregel_send?.([[stateKey, evt]]);
     return evt;
   };
 

libs/sdk-py/langgraph_sdk/client.py

@@ -1831,7 +1831,12 @@ class RunsClient:
         return await self.http.get(f"/threads/{thread_id}/runs/{run_id}/join")
 
     def join_stream(
-        self, thread_id: str, run_id: str, *, cancel_on_disconnect: bool = False
+        self,
+        thread_id: str,
+        run_id: str,
+        *,
+        cancel_on_disconnect: bool = False,
+        stream_mode: Optional[Union[StreamMode, Sequence[StreamMode]]] = None,
     ) -> AsyncIterator[StreamPart]:
         """Stream output from a run in real-time, until the run is done.
         Output is not buffered, so any output produced before this call will
@@ -1841,6 +1846,9 @@ class RunsClient:
             thread_id: The thread ID to join.
             run_id: The run ID to join.
             cancel_on_disconnect: Whether to cancel the run when the stream is disconnected.
+            stream_mode: The stream mode(s) to use. Must be a subset of the stream modes passed
+                when creating the run. Background runs default to having the union of all
+                stream modes.
 
         Returns:
             None
@@ -1849,14 +1857,18 @@ class RunsClient:
 
             await client.runs.join_stream(
                 thread_id="thread_id_to_join",
-                run_id="run_id_to_join"
+                run_id="run_id_to_join",
+                stream_mode=["values", "debug"]
             )
 
         """  # noqa: E501
         return self.http.stream(
             f"/threads/{thread_id}/runs/{run_id}/stream",
             "GET",
-            params={"cancel_on_disconnect": cancel_on_disconnect},
+            params={
+                "cancel_on_disconnect": cancel_on_disconnect,
+                "stream_mode": stream_mode,
+            },
         )
 
     async def delete(self, thread_id: str, run_id: str) -> None:
@@ -3988,7 +4000,14 @@ class SyncRunsClient:
         """  # noqa: E501
         return self.http.get(f"/threads/{thread_id}/runs/{run_id}/join")
 
-    def join_stream(self, thread_id: str, run_id: str) -> Iterator[StreamPart]:
+    def join_stream(
+        self,
+        thread_id: str,
+        run_id: str,
+        *,
+        stream_mode: Optional[Union[StreamMode, Sequence[StreamMode]]] = None,
+        cancel_on_disconnect: bool = False,
+    ) -> Iterator[StreamPart]:
         """Stream output from a run in real-time, until the run is done.
         Output is not buffered, so any output produced before this call will
         not be received here.
@@ -3996,6 +4015,10 @@ class SyncRunsClient:
         Args:
             thread_id: The thread ID to join.
             run_id: The run ID to join.
+            stream_mode: The stream mode(s) to use. Must be a subset of the stream modes passed
+                when creating the run. Background runs default to having the union of all
+                stream modes.
+            cancel_on_disconnect: Whether to cancel the run when the stream is disconnected.
 
         Returns:
             None
@@ -4004,11 +4027,19 @@ class SyncRunsClient:
 
             client.runs.join_stream(
                 thread_id="thread_id_to_join",
-                run_id="run_id_to_join"
+                run_id="run_id_to_join",
+                stream_mode=["values", "debug"]
             )
 
         """  # noqa: E501
-        return self.http.stream(f"/threads/{thread_id}/runs/{run_id}/stream", "GET")
+        return self.http.stream(
+            f"/threads/{thread_id}/runs/{run_id}/stream",
+            "GET",
+            params={
+                "stream_mode": stream_mode,
+                "cancel_on_disconnect": cancel_on_disconnect,
+            },
+        )
 
     def delete(self, thread_id: str, run_id: str) -> None:
         """Delete a run.

libs/sdk-py/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph-sdk"
-version = "0.1.56"
+version = "0.1.57"
 description = "SDK for interacting with LangGraph API"
 authors = []
 license = "MIT"