mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-24 00:22:25 +02:00
Merge branch 'main' into pyupgrade-39
This commit is contained in:
Sydney Runkle
@@ -1,7 +1,7 @@
|
||||
# LangGraph Studio With Local Deployment
|
||||
|
||||
!!! warning "Browser Compatibility"
|
||||
Viewing the studio page of a local LangGraph deployment does not work in Safari. Use Chrome instead.
|
||||
Safari blocks `localhost` connections to Studio. To work around this, start the server with `--tunnel` and you’ll be able to access Studio from Safari via a secure tunnel.
|
||||
|
||||
## Setup
|
||||
|
||||
|
||||
@@ -10,9 +10,6 @@ The LangGraph command line interface includes commands to build and run a LangGr
|
||||
=== "Python"
|
||||
```bash
|
||||
pip install langgraph-cli
|
||||
|
||||
# Install via Homebrew
|
||||
brew install langgraph-cli
|
||||
```
|
||||
|
||||
=== "JS"
|
||||
@@ -298,6 +295,11 @@ The LangGraph CLI requires a JSON configuration file that follows this [schema](
|
||||
| `--no-reload` | | Disable auto-reload |
|
||||
| `--n-jobs-per-worker INTEGER` | | Number of jobs per worker. Default is 10 |
|
||||
| `--debug-port INTEGER` | | Port for debugger to listen on |
|
||||
| `--wait-for-client` | `False` | Wait for a debugger client to connect to the debug port before starting the server |
|
||||
| `--no-browser` | | Skip automatically opening the browser when the server starts |
|
||||
| `--studio-url TEXT` | | URL of the LangGraph Studio instance to connect to. Defaults to https://smith.langchain.com |
|
||||
| `--allow-blocking` | `False` | Do not raise errors for synchronous I/O blocking operations in your code (added in `0.2.6`) |
|
||||
| `--tunnel` | `False` | Expose the local server via a public tunnel (Cloudflare) for remote frontend access. This avoids issues with browsers like Safari or networks blocking localhost connections |
|
||||
| `--help` | | Display command documentation |
|
||||
|
||||
|
||||
@@ -321,6 +323,11 @@ The LangGraph CLI requires a JSON configuration file that follows this [schema](
|
||||
| `--no-reload` | | Disable auto-reload |
|
||||
| `--n-jobs-per-worker INTEGER` | | Number of jobs per worker. Default is 10 |
|
||||
| `--debug-port INTEGER` | | Port for debugger to listen on |
|
||||
| `--wait-for-client` | `False` | Wait for a debugger client to connect to the debug port before starting the server |
|
||||
| `--no-browser` | | Skip automatically opening the browser when the server starts |
|
||||
| `--studio-url TEXT` | | URL of the LangGraph Studio instance to connect to. Defaults to https://smith.langchain.com |
|
||||
| `--allow-blocking` | `False` | Do not raise errors for synchronous I/O blocking operations in your code |
|
||||
| `--tunnel` | `False` | Expose the local server via a public tunnel (Cloudflare) for remote frontend access. This avoids issues with browsers or networks blocking localhost connections |
|
||||
| `--help` | | Display command documentation |
|
||||
|
||||
### `build`
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
.safari {
|
||||
color: #0070C9;
|
||||
}
|
||||
@@ -14,3 +14,4 @@ Errors referenced below will have an `lc_error_code` property corresponding to o
|
||||
These guides provide troubleshooting information for errors that are specific to the LangGraph Platform.
|
||||
|
||||
- [INVALID_LICENSE](./INVALID_LICENSE.md)
|
||||
- [Studio Errors](../studio.md)
|
||||
|
||||
@@ -0,0 +1,45 @@
|
||||
# Troubleshooting LangGraph Studio
|
||||
|
||||
## :fontawesome-brands-safari:{ .safari } Safari connection error with local dev server
|
||||
|
||||
Safari blocks plain‑HTTP traffic on localhost. If you start Studio with a vanilla
|
||||
`langgraph dev`, the page may report a "Failed to load assistants" error (or something similar) and the browser DevTools will show network errors.
|
||||
|
||||
#### Quick fix — run Studio through a secure Cloudflare tunnel
|
||||
|
||||
=== "Python"
|
||||
|
||||
```shell
|
||||
pip install -U langgraph-cli>=0.2.6 # Python
|
||||
langgraph dev --tunnel
|
||||
```
|
||||
=== "JS"
|
||||
|
||||
```shell
|
||||
# Requires @langchain/langgraph-cli>=0.0.26
|
||||
npx @langchain/langgraph-cli dev
|
||||
```
|
||||
|
||||
The command prints a URL like:
|
||||
|
||||
```shell
|
||||
https://smith.langchain.com/studio/?baseUrl=https://hamilton-praise-heart-costumes.trycloudflare.com
|
||||
```
|
||||
where
|
||||
```shell
|
||||
?baseUrl=https://hamilton-praise-heart-costumes.trycloudflare.com
|
||||
```
|
||||
indicates the endpoint where your agent server is exposed.
|
||||
|
||||
Open that URL in Safari and Studio should load immediately.
|
||||
|
||||
#### Alternative — use a Chromium‑based browser
|
||||
|
||||
Chrome, Edge, and Brave allow HTTP on localhost, so a plain `langgraph dev` should work without extra steps.
|
||||
|
||||
#### If it’s still not loading
|
||||
|
||||
1. Make sure the `baseUrl` query parameter in the studio URL points to the **tunnel URL** NOT to localhost.
|
||||
2. Confirm your CLI version with `langgraph --version`.
|
||||
|
||||
No other configuration, certificates, or CORS tweaks are required.
|
||||
+4
-2
@@ -1,5 +1,5 @@
|
||||
site_name: ""
|
||||
site_description: Build language agents as graphs
|
||||
site_name: "LangGraph"
|
||||
site_description: Build reliable, stateful AI systems, without giving up control
|
||||
site_url: https://langchain-ai.github.io/langgraph/
|
||||
repo_url: https://github.com/langchain-ai/langgraph
|
||||
edit_uri: edit/main/docs/docs/
|
||||
@@ -400,6 +400,7 @@ nav:
|
||||
- troubleshooting/errors/MULTIPLE_SUBGRAPHS.md
|
||||
- troubleshooting/errors/INVALID_CHAT_HISTORY.md
|
||||
- troubleshooting/errors/INVALID_LICENSE.md
|
||||
- troubleshooting/studio.md
|
||||
- LangGraph Academy Course: https://academy.langchain.com/courses/intro-to-langgraph
|
||||
|
||||
- Agents:
|
||||
@@ -549,3 +550,4 @@ copyright: >
|
||||
Copyright © 2025 LangChain, Inc | <a href="#__consent">Consent Preferences</a>
|
||||
extra_css:
|
||||
- stylesheets/version_admonitions.css
|
||||
- stylesheets/logos.css
|
||||
|
||||
@@ -572,6 +572,14 @@ def dockerfile(save_path: str, config: pathlib.Path, add_docker_compose: bool) -
|
||||
help="Don't raise errors for synchronous I/O blocking operations in your code.",
|
||||
default=False,
|
||||
)
|
||||
@click.option(
|
||||
"--tunnel",
|
||||
is_flag=True,
|
||||
help="Expose the local server via a public tunnel (in this case, Cloudflare) "
|
||||
"for remote frontend access. This avoids issues with browsers "
|
||||
"or networks blocking localhost connections.",
|
||||
default=False,
|
||||
)
|
||||
@cli.command(
|
||||
"dev",
|
||||
help="🏃♀️➡️ Run LangGraph API server in development mode with hot reloading and debugging support",
|
||||
@@ -588,6 +596,7 @@ def dev(
|
||||
wait_for_client: bool,
|
||||
studio_url: Optional[str],
|
||||
allow_blocking: bool,
|
||||
tunnel: bool,
|
||||
):
|
||||
"""CLI entrypoint for running the LangGraph API server."""
|
||||
try:
|
||||
@@ -655,6 +664,7 @@ def dev(
|
||||
ui_config=config_json.get("ui_config"),
|
||||
studio_url=studio_url,
|
||||
allow_blocking=allow_blocking,
|
||||
tunnel=tunnel,
|
||||
)
|
||||
|
||||
|
||||
|
||||
Generated
+16
-16
@@ -585,15 +585,15 @@ tests = ["flask (>=2.2.5)", "hypothesis (>=6.79.4)", "pytest (>=7.4.4)"]
|
||||
|
||||
[[package]]
|
||||
name = "langchain-core"
|
||||
version = "0.3.54"
|
||||
version = "0.3.55"
|
||||
description = "Building applications with LLMs through composability"
|
||||
optional = true
|
||||
python-versions = "<4.0,>=3.9"
|
||||
groups = ["main"]
|
||||
markers = "python_version >= \"3.11\" and extra == \"inmem\""
|
||||
files = [
|
||||
{file = "langchain_core-0.3.54-py3-none-any.whl", hash = "sha256:cd42155d9089e2fd4695ee02a4b2bc6daf55b9d4e1a37639647cf2455ed4fa04"},
|
||||
{file = "langchain_core-0.3.54.tar.gz", hash = "sha256:55ce38939038e19b1271f36f512335462d7f64057b531598b3651d2b403e1b42"},
|
||||
{file = "langchain_core-0.3.55-py3-none-any.whl", hash = "sha256:b3cb36bf37755a616158a79866657c6697b43a2f7c69dd723ce425f1c76c1baa"},
|
||||
{file = "langchain_core-0.3.55.tar.gz", hash = "sha256:0f2b3e311621116a83510c70b0ac9d959030a0a457a69483535cff18501fedc9"},
|
||||
]
|
||||
|
||||
[package.dependencies]
|
||||
@@ -630,15 +630,15 @@ xxhash = ">=3.5.0,<4.0.0"
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-api"
|
||||
version = "0.1.9"
|
||||
version = "0.1.12"
|
||||
description = ""
|
||||
optional = true
|
||||
python-versions = "<4.0,>=3.11.0"
|
||||
groups = ["main"]
|
||||
markers = "python_version >= \"3.11\" and extra == \"inmem\""
|
||||
files = [
|
||||
{file = "langgraph_api-0.1.9-py3-none-any.whl", hash = "sha256:f84b11b1855e68dbef9f0a78db803e325b8dc11e2e19613178ab84cb0d99627c"},
|
||||
{file = "langgraph_api-0.1.9.tar.gz", hash = "sha256:3530d82e715b9f99eeb8753c365f4d16c99ce60533fa15530d5ad1d493aeec06"},
|
||||
{file = "langgraph_api-0.1.12-py3-none-any.whl", hash = "sha256:0f9417052ac75f6da892902083b7cf6a515bee12dd035cfbc3f0ddb813738830"},
|
||||
{file = "langgraph_api-0.1.12.tar.gz", hash = "sha256:1646a904121a5dc84cece6a81b9c49693ccfbd6f1a2904e0ed7aa7eb711e64fc"},
|
||||
]
|
||||
|
||||
[package.dependencies]
|
||||
@@ -712,15 +712,15 @@ blockbuster = ">=1.5.24,<2.0.0"
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-sdk"
|
||||
version = "0.1.61"
|
||||
version = "0.1.63"
|
||||
description = "SDK for interacting with LangGraph API"
|
||||
optional = true
|
||||
python-versions = "<4.0.0,>=3.9.0"
|
||||
groups = ["main"]
|
||||
markers = "python_version >= \"3.11\""
|
||||
files = [
|
||||
{file = "langgraph_sdk-0.1.61-py3-none-any.whl", hash = "sha256:f2d774b12497c428862993090622d51e0dbc3f53e0cee3d74a13c7495d835cc6"},
|
||||
{file = "langgraph_sdk-0.1.61.tar.gz", hash = "sha256:87dd1f07ab82da8875ac343268ece8bf5414632017ebc9d1cef4b523962fd601"},
|
||||
{file = "langgraph_sdk-0.1.63-py3-none-any.whl", hash = "sha256:6fb78a7fc6a30eea43bd0d6401dbc9e3263d0d4c03f63c04035980da7e586b05"},
|
||||
{file = "langgraph_sdk-0.1.63.tar.gz", hash = "sha256:62bf2cc31e5aa6c5b9011ee1702bcf1e36e67e142a60bd97af2611162fb58e18"},
|
||||
]
|
||||
|
||||
[package.dependencies]
|
||||
@@ -729,15 +729,15 @@ orjson = ">=3.10.1"
|
||||
|
||||
[[package]]
|
||||
name = "langsmith"
|
||||
version = "0.3.32"
|
||||
version = "0.3.33"
|
||||
description = "Client library to connect to the LangSmith LLM Tracing and Evaluation Platform."
|
||||
optional = true
|
||||
python-versions = "<4.0,>=3.9"
|
||||
groups = ["main"]
|
||||
markers = "python_version >= \"3.11\" and extra == \"inmem\""
|
||||
files = [
|
||||
{file = "langsmith-0.3.32-py3-none-any.whl", hash = "sha256:d79af299038cd13db6d53f99fdc6a7171b731702536b7635f2da132555a6bebc"},
|
||||
{file = "langsmith-0.3.32.tar.gz", hash = "sha256:3d7b1149e9fbe0f388303bc94d8deeeb822acc89cdea34f24151ea1316eb487d"},
|
||||
{file = "langsmith-0.3.33-py3-none-any.whl", hash = "sha256:6fa453942014945e1de7e283880ed3e8031b5d84e0dc75b87d101ecedb62371b"},
|
||||
{file = "langsmith-0.3.33.tar.gz", hash = "sha256:0f439e945528c6d14140137b918cc048aea04c6a987525926dbfda2560002924"},
|
||||
]
|
||||
|
||||
[package.dependencies]
|
||||
@@ -1611,15 +1611,15 @@ zstd = ["zstandard (>=0.18.0)"]
|
||||
|
||||
[[package]]
|
||||
name = "uvicorn"
|
||||
version = "0.34.1"
|
||||
version = "0.34.2"
|
||||
description = "The lightning-fast ASGI server."
|
||||
optional = true
|
||||
python-versions = ">=3.9"
|
||||
groups = ["main"]
|
||||
markers = "python_version >= \"3.11\" and extra == \"inmem\""
|
||||
files = [
|
||||
{file = "uvicorn-0.34.1-py3-none-any.whl", hash = "sha256:984c3a8c7ca18ebaad15995ee7401179212c59521e67bfc390c07fa2b8d2e065"},
|
||||
{file = "uvicorn-0.34.1.tar.gz", hash = "sha256:af981725fc4b7ffc5cb3b0e9eda6258a90c4b52cb2a83ce567ae0a7ae1757afc"},
|
||||
{file = "uvicorn-0.34.2-py3-none-any.whl", hash = "sha256:deb49af569084536d269fe0a6d67e3754f104cf03aba7c11c40f01aadf33c403"},
|
||||
{file = "uvicorn-0.34.2.tar.gz", hash = "sha256:0e929828f6186353a80b58ea719861d2629d766293b6d19baf086ba31d4f3328"},
|
||||
]
|
||||
|
||||
[package.dependencies]
|
||||
@@ -2011,4 +2011,4 @@ inmem = ["langgraph-api", "langgraph-runtime-inmem", "python-dotenv"]
|
||||
[metadata]
|
||||
lock-version = "2.1"
|
||||
python-versions = "^3.9.0,<4.0"
|
||||
content-hash = "afc2f8776b4b6144bd1197df49ba34089889e2a1110b8470d8f1b212e0b08380"
|
||||
content-hash = "6f3f275ae70749922db5bd1105711fbb0f8ac3b215982f7c75861c93d956e95d"
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[tool.poetry]
|
||||
name = "langgraph-cli"
|
||||
version = "0.2.5"
|
||||
version = "0.2.6"
|
||||
description = "CLI for interacting with LangGraph API"
|
||||
authors = []
|
||||
license = "MIT"
|
||||
@@ -14,7 +14,7 @@ langgraph = "langgraph_cli.cli:cli"
|
||||
[tool.poetry.dependencies]
|
||||
python = "^3.9.0,<4.0"
|
||||
click = "^8.1.7"
|
||||
langgraph-api = { version = ">=0.1.0,<0.2.0", optional = true, python = ">=3.11,<4.0" }
|
||||
langgraph-api = { version = ">=0.1.12,<0.2.0", optional = true, python = ">=3.11,<4.0" }
|
||||
langgraph-runtime-inmem = { version = ">=0.0.1,<0.1.0", optional = true, python = ">=3.11,<4.0" }
|
||||
langgraph-sdk = { version = ">=0.1.0,<0.2.0", optional = true, python = ">=3.11,<4.0" }
|
||||
python-dotenv = { version = ">=0.8.0", optional = true }
|
||||
|
||||
@@ -4,6 +4,7 @@ from inspect import (
|
||||
ismethod,
|
||||
signature,
|
||||
)
|
||||
from itertools import zip_longest
|
||||
from types import FunctionType
|
||||
from typing import (
|
||||
Any,
|
||||
@@ -26,12 +27,17 @@ from langchain_core.runnables import (
|
||||
|
||||
from langgraph.constants import END, START
|
||||
from langgraph.errors import InvalidUpdateError
|
||||
from langgraph.pregel.write import ChannelWrite
|
||||
from langgraph.pregel.write import PASSTHROUGH, ChannelWrite, ChannelWriteEntry
|
||||
from langgraph.types import Send
|
||||
from langgraph.utils.runnable import (
|
||||
RunnableCallable,
|
||||
)
|
||||
|
||||
Writer = Callable[
|
||||
[Sequence[Union[str, Send]]],
|
||||
Sequence[Union[ChannelWriteEntry, Send]],
|
||||
]
|
||||
|
||||
|
||||
def _get_branch_path_input_schema(
|
||||
path: Union[
|
||||
@@ -121,9 +127,7 @@ class Branch(NamedTuple):
|
||||
|
||||
def run(
|
||||
self,
|
||||
writer: Callable[
|
||||
[Sequence[Union[str, Send]], RunnableConfig], Optional[ChannelWrite]
|
||||
],
|
||||
writer: Writer,
|
||||
reader: Optional[Callable[[RunnableConfig], Any]] = None,
|
||||
) -> RunnableCallable:
|
||||
return ChannelWrite.register_writer(
|
||||
@@ -135,7 +139,15 @@ class Branch(NamedTuple):
|
||||
name=None,
|
||||
trace=False,
|
||||
func_accepts_config=True,
|
||||
),
|
||||
list(
|
||||
zip_longest(
|
||||
writer([e for e in self.ends.values() if e != END]),
|
||||
[str(la) for la, e in self.ends.items() if e != END],
|
||||
)
|
||||
)
|
||||
if self.ends
|
||||
else None,
|
||||
)
|
||||
|
||||
def _route(
|
||||
@@ -144,9 +156,7 @@ class Branch(NamedTuple):
|
||||
config: RunnableConfig,
|
||||
*,
|
||||
reader: Optional[Callable[[RunnableConfig], Any]],
|
||||
writer: Callable[
|
||||
[Sequence[Union[str, Send]], RunnableConfig], Optional[ChannelWrite]
|
||||
],
|
||||
writer: Writer,
|
||||
) -> Runnable:
|
||||
if reader:
|
||||
value = reader(config)
|
||||
@@ -169,9 +179,7 @@ class Branch(NamedTuple):
|
||||
config: RunnableConfig,
|
||||
*,
|
||||
reader: Optional[Callable[[RunnableConfig], Any]],
|
||||
writer: Callable[
|
||||
[Sequence[Union[str, Send]], RunnableConfig], Optional[ChannelWrite]
|
||||
],
|
||||
writer: Writer,
|
||||
) -> Runnable:
|
||||
if reader:
|
||||
value = reader(config)
|
||||
@@ -190,9 +198,7 @@ class Branch(NamedTuple):
|
||||
|
||||
def _finish(
|
||||
self,
|
||||
writer: Callable[
|
||||
[Sequence[Union[str, Send]], RunnableConfig], Optional[ChannelWrite]
|
||||
],
|
||||
writer: Writer,
|
||||
input: Any,
|
||||
result: Any,
|
||||
config: RunnableConfig,
|
||||
@@ -209,4 +215,18 @@ class Branch(NamedTuple):
|
||||
raise ValueError("Branch did not return a valid destination")
|
||||
if any(p.node == END for p in destinations if isinstance(p, Send)):
|
||||
raise InvalidUpdateError("Cannot send a packet to the END node")
|
||||
return writer(destinations, config) or input
|
||||
entries = writer(destinations)
|
||||
if not entries:
|
||||
return input
|
||||
else:
|
||||
need_passthrough = False
|
||||
for e in entries:
|
||||
if isinstance(e, ChannelWriteEntry):
|
||||
if e.value is PASSTHROUGH:
|
||||
need_passthrough = True
|
||||
break
|
||||
if need_passthrough:
|
||||
return ChannelWrite(entries)
|
||||
else:
|
||||
ChannelWrite.do_write(config, entries)
|
||||
return input
|
||||
|
||||
@@ -1,4 +1,3 @@
|
||||
import asyncio
|
||||
import logging
|
||||
from collections import defaultdict
|
||||
from collections.abc import Awaitable, Hashable, Sequence
|
||||
@@ -13,9 +12,6 @@ from typing import (
|
||||
)
|
||||
|
||||
from langchain_core.runnables import Runnable
|
||||
from langchain_core.runnables.config import RunnableConfig
|
||||
from langchain_core.runnables.graph import Graph as DrawableGraph
|
||||
from langchain_core.runnables.graph import Node as DrawableNode
|
||||
from typing_extensions import Self
|
||||
|
||||
from langgraph.channels.ephemeral_value import EphemeralValue
|
||||
@@ -30,7 +26,6 @@ from langgraph.constants import (
|
||||
)
|
||||
from langgraph.graph.branch import Branch
|
||||
from langgraph.pregel import Channel, Pregel
|
||||
from langgraph.pregel.protocol import PregelProtocol
|
||||
from langgraph.pregel.read import PregelNode
|
||||
from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
|
||||
from langgraph.types import All, Checkpointer
|
||||
@@ -378,10 +373,10 @@ class CompiledGraph(Pregel):
|
||||
cast(list[str], self.nodes[end].channels).append(start)
|
||||
|
||||
def attach_branch(self, start: str, name: str, branch: Branch) -> None:
|
||||
def branch_writer(
|
||||
packets: Sequence[Union[str, Send]], config: RunnableConfig
|
||||
) -> Optional[ChannelWrite]:
|
||||
writes = [
|
||||
def get_writes(
|
||||
packets: Sequence[Union[str, Send]],
|
||||
) -> Sequence[Union[ChannelWriteEntry, Send]]:
|
||||
return [
|
||||
(
|
||||
ChannelWriteEntry(f"branch:{start}:{name}:{p}" if p != END else END)
|
||||
if not isinstance(p, Send)
|
||||
@@ -389,14 +384,13 @@ class CompiledGraph(Pregel):
|
||||
)
|
||||
for p in packets
|
||||
]
|
||||
return ChannelWrite(cast(Sequence[Union[ChannelWriteEntry, Send]], writes))
|
||||
|
||||
# add hidden start node
|
||||
if start == START and start not in self.nodes:
|
||||
self.nodes[start] = Channel.subscribe_to(START, tags=[TAG_HIDDEN])
|
||||
|
||||
# attach branch writer
|
||||
self.nodes[start] |= branch.run(branch_writer)
|
||||
self.nodes[start] |= branch.run(get_writes)
|
||||
|
||||
# attach branch readers
|
||||
ends = branch.ends.values() if branch.ends else [node for node in self.nodes]
|
||||
@@ -406,171 +400,3 @@ class CompiledGraph(Pregel):
|
||||
self.channels[channel_name] = EphemeralValue(Any)
|
||||
self.nodes[end].triggers.append(channel_name)
|
||||
cast(list[str], self.nodes[end].channels).append(channel_name)
|
||||
|
||||
async def aget_graph(
|
||||
self,
|
||||
config: Optional[RunnableConfig] = None,
|
||||
*,
|
||||
xray: Union[int, bool] = False,
|
||||
) -> DrawableGraph:
|
||||
"""Returns a drawable representation of the computation graph."""
|
||||
from langgraph.pregel.remote import RemoteGraph
|
||||
|
||||
# gather subgraphs
|
||||
if xray:
|
||||
subpregels: dict[str, PregelProtocol] = {
|
||||
k: v
|
||||
async for k, v in self.aget_subgraphs()
|
||||
if isinstance(v, (CompiledGraph, RemoteGraph))
|
||||
}
|
||||
subgraphs = {
|
||||
k: v
|
||||
for k, v in zip(
|
||||
subpregels,
|
||||
await asyncio.gather(
|
||||
*(
|
||||
p.aget_graph(
|
||||
config,
|
||||
xray=xray
|
||||
if isinstance(xray, bool) or xray <= 0
|
||||
else xray - 1,
|
||||
)
|
||||
for p in subpregels.values()
|
||||
)
|
||||
),
|
||||
)
|
||||
}
|
||||
else:
|
||||
subgraphs = {}
|
||||
|
||||
# draw the graph
|
||||
return self._draw_graph(config, subgraphs=subgraphs)
|
||||
|
||||
def get_graph(
|
||||
self,
|
||||
config: Optional[RunnableConfig] = None,
|
||||
*,
|
||||
xray: Union[int, bool] = False,
|
||||
) -> DrawableGraph:
|
||||
"""Returns a drawable representation of the computation graph."""
|
||||
from langgraph.pregel.remote import RemoteGraph
|
||||
|
||||
# gather subgraphs
|
||||
if xray:
|
||||
subgraphs = {
|
||||
k: v.get_graph(
|
||||
config,
|
||||
xray=xray if isinstance(xray, bool) or xray <= 0 else xray - 1,
|
||||
)
|
||||
for k, v in self.get_subgraphs()
|
||||
if isinstance(v, (CompiledGraph, RemoteGraph))
|
||||
}
|
||||
else:
|
||||
subgraphs = {}
|
||||
|
||||
# draw the graph
|
||||
return self._draw_graph(config, subgraphs=subgraphs)
|
||||
|
||||
def _draw_graph(
|
||||
self,
|
||||
config: Optional[RunnableConfig] = None,
|
||||
*,
|
||||
subgraphs: dict[str, DrawableGraph] = {},
|
||||
) -> DrawableGraph:
|
||||
# create the graph
|
||||
graph = DrawableGraph()
|
||||
start_nodes: dict[str, DrawableNode] = {
|
||||
START: graph.add_node(self.get_input_schema(config), START)
|
||||
}
|
||||
end_nodes: dict[str, DrawableNode] = {}
|
||||
|
||||
def add_edge(
|
||||
start: str,
|
||||
end: str,
|
||||
label: Optional[Hashable] = None,
|
||||
conditional: bool = False,
|
||||
) -> None:
|
||||
if end == END and END not in end_nodes:
|
||||
end_nodes[END] = graph.add_node(self.get_output_schema(config), END)
|
||||
if start not in start_nodes or end not in end_nodes:
|
||||
logger.warning(
|
||||
f"Could not add edge from '{start}' to '{end}' due to missing nodes"
|
||||
)
|
||||
return
|
||||
return graph.add_edge(
|
||||
start_nodes[start],
|
||||
end_nodes[end],
|
||||
str(label) if label is not None else None,
|
||||
conditional,
|
||||
)
|
||||
|
||||
for key, n in self.builder.nodes.items():
|
||||
node = n.runnable
|
||||
metadata = n.metadata or {}
|
||||
if key in self.interrupt_before_nodes and key in self.interrupt_after_nodes:
|
||||
metadata["__interrupt"] = "before,after"
|
||||
elif key in self.interrupt_before_nodes:
|
||||
metadata["__interrupt"] = "before"
|
||||
elif key in self.interrupt_after_nodes:
|
||||
metadata["__interrupt"] = "after"
|
||||
if key in subgraphs:
|
||||
subgraph = subgraphs[key]
|
||||
subgraph.trim_first_node()
|
||||
subgraph.trim_last_node()
|
||||
if len(subgraph.nodes) >= 1:
|
||||
e, s = graph.extend(subgraph, prefix=key)
|
||||
if e is None:
|
||||
logger.warning(
|
||||
f"Could not extend subgraph '{key}' due to missing entrypoint"
|
||||
)
|
||||
continue
|
||||
if s is not None:
|
||||
start_nodes[key] = s
|
||||
end_nodes[key] = e
|
||||
else:
|
||||
nn = graph.add_node(node, key, metadata=metadata or None)
|
||||
start_nodes[key] = nn
|
||||
end_nodes[key] = nn
|
||||
else:
|
||||
nn = graph.add_node(node, key, metadata=metadata or None)
|
||||
start_nodes[key] = nn
|
||||
end_nodes[key] = nn
|
||||
for start, end in sorted(self.builder._all_edges):
|
||||
add_edge(start, end)
|
||||
for start, branches in self.builder.branches.items():
|
||||
default_ends = {
|
||||
**{k: k for k in self.builder.nodes if k != start},
|
||||
END: END,
|
||||
}
|
||||
for _, branch in branches.items():
|
||||
if branch.ends is not None:
|
||||
ends = branch.ends
|
||||
elif branch.then is not None:
|
||||
ends = {k: k for k in default_ends if k not in (END, branch.then)}
|
||||
else:
|
||||
ends = cast(dict[Hashable, str], default_ends)
|
||||
for label, end in ends.items():
|
||||
add_edge(
|
||||
start,
|
||||
end,
|
||||
label if label != end else None,
|
||||
conditional=True,
|
||||
)
|
||||
if branch.then is not None:
|
||||
add_edge(end, branch.then)
|
||||
for key, n in self.builder.nodes.items():
|
||||
if isinstance(n.ends, dict):
|
||||
for end, label in n.ends.items():
|
||||
add_edge(key, end, label, conditional=True)
|
||||
elif isinstance(n.ends, tuple):
|
||||
for end in n.ends:
|
||||
add_edge(key, end, conditional=True)
|
||||
|
||||
return graph
|
||||
|
||||
def _repr_mimebundle_(self, **kwargs: Any) -> dict[str, Any]:
|
||||
"""Mime bundle used by Jupyter to display the graph"""
|
||||
return {
|
||||
"text/plain": repr(self),
|
||||
"image/png": self.get_graph().draw_mermaid_png(),
|
||||
}
|
||||
|
||||
@@ -771,7 +771,12 @@ class CompiledStateGraph(CompiledGraph):
|
||||
ChannelWriteTupleEntry(
|
||||
mapper=_get_root if output_keys == ["__root__"] else _get_updates
|
||||
),
|
||||
ChannelWriteTupleEntry(mapper=_control_branch),
|
||||
ChannelWriteTupleEntry(
|
||||
mapper=_control_branch,
|
||||
static=_control_static(node.ends)
|
||||
if node is not None and node.ends is not None
|
||||
else None,
|
||||
),
|
||||
)
|
||||
|
||||
# add node and output channel
|
||||
@@ -837,9 +842,9 @@ class CompiledStateGraph(CompiledGraph):
|
||||
def attach_branch(
|
||||
self, start: str, name: str, branch: Branch, *, with_reader: bool = True
|
||||
) -> None:
|
||||
def branch_writer(
|
||||
packets: Sequence[Union[str, Send]], config: RunnableConfig
|
||||
) -> None:
|
||||
def get_writes(
|
||||
packets: Sequence[Union[str, Send]],
|
||||
) -> Sequence[Union[ChannelWriteEntry, Send]]:
|
||||
if filtered := [p for p in packets if p != END]:
|
||||
writes = [
|
||||
(
|
||||
@@ -854,13 +859,15 @@ class CompiledStateGraph(CompiledGraph):
|
||||
ChannelWriteEntry(
|
||||
f"branch:{start}:{name}::then",
|
||||
WaitForNames(
|
||||
{p.node if isinstance(p, Send) else p for p in filtered}
|
||||
frozenset(
|
||||
p.node if isinstance(p, Send) else p
|
||||
for p in filtered
|
||||
)
|
||||
),
|
||||
)
|
||||
)
|
||||
ChannelWrite.do_write(
|
||||
config, cast(Sequence[Union[Send, ChannelWriteEntry]], writes)
|
||||
)
|
||||
return writes
|
||||
return []
|
||||
|
||||
if with_reader:
|
||||
# get schema
|
||||
@@ -888,7 +895,7 @@ class CompiledStateGraph(CompiledGraph):
|
||||
reader = None
|
||||
|
||||
# attach branch publisher
|
||||
self.nodes[start].writers.append(branch.run(branch_writer, reader))
|
||||
self.nodes[start].writers.append(branch.run(get_writes, reader))
|
||||
|
||||
# attach then subscriber
|
||||
if branch.then and branch.then != END:
|
||||
@@ -1056,6 +1063,19 @@ def _control_branch(value: Any) -> Sequence[tuple[str, Any]]:
|
||||
return rtn
|
||||
|
||||
|
||||
def _control_static(
|
||||
ends: Union[tuple[str, ...], dict[str, str]],
|
||||
) -> Sequence[tuple[str, Any, Optional[str]]]:
|
||||
if isinstance(ends, dict):
|
||||
return [
|
||||
(CHANNEL_BRANCH_TO.format(k), None, label)
|
||||
for k, label in ends.items()
|
||||
if k != END
|
||||
]
|
||||
else:
|
||||
return [(CHANNEL_BRANCH_TO.format(e), None, None) for e in ends if e != END]
|
||||
|
||||
|
||||
def _get_root(input: Any) -> Optional[Sequence[tuple[str, Any]]]:
|
||||
if isinstance(input, Command):
|
||||
if input.graph == Command.PARENT:
|
||||
|
||||
@@ -87,6 +87,7 @@ from langgraph.pregel.algo import (
|
||||
)
|
||||
from langgraph.pregel.checkpoint import create_checkpoint, empty_checkpoint
|
||||
from langgraph.pregel.debug import tasks_w_writes
|
||||
from langgraph.pregel.draw import draw_graph
|
||||
from langgraph.pregel.io import map_input, read_channels
|
||||
from langgraph.pregel.loop import AsyncPregelLoop, StreamProtocol, SyncPregelLoop
|
||||
from langgraph.pregel.manager import AsyncChannelsManager, ChannelsManager
|
||||
@@ -556,14 +557,84 @@ class Pregel(PregelProtocol):
|
||||
self.validate()
|
||||
|
||||
def get_graph(
|
||||
self, config: RunnableConfig | None = None, *, xray: int | bool = False
|
||||
self, config: RunnableConfig | None = None, *, xray: int | bool = Fals
|
||||
) -> Graph:
|
||||
raise NotImplementedError
|
||||
"""Returns a drawable representation of the computation graph."""
|
||||
# gather subgraphs
|
||||
if xray:
|
||||
subgraphs = {
|
||||
k: v.get_graph(
|
||||
config,
|
||||
xray=xray if isinstance(xray, bool) or xray <= 0 else xray - 1,
|
||||
)
|
||||
for k, v in self.get_subgraphs()
|
||||
}
|
||||
else:
|
||||
subgraphs = {}
|
||||
|
||||
return draw_graph(
|
||||
merge_configs(self.config, config),
|
||||
nodes=self.nodes,
|
||||
specs=self.channels,
|
||||
input_channels=self.input_channels,
|
||||
interrupt_after_nodes=self.interrupt_after_nodes,
|
||||
interrupt_before_nodes=self.interrupt_before_nodes,
|
||||
trigger_to_nodes=self.trigger_to_nodes,
|
||||
checkpointer=self.checkpointer,
|
||||
subgraphs=subgraphs,
|
||||
)
|
||||
|
||||
async def aget_graph(
|
||||
self, config: RunnableConfig | None = None, *, xray: int | bool = False
|
||||
self,
|
||||
config: RunnableConfig | None = None,
|
||||
*,
|
||||
xray: int | bool = False
|
||||
) -> Graph:
|
||||
raise NotImplementedError
|
||||
"""Returns a drawable representation of the computation graph."""
|
||||
|
||||
# gather subgraphs
|
||||
if xray:
|
||||
subpregels: dict[str, PregelProtocol] = {
|
||||
k: v async for k, v in self.aget_subgraphs()
|
||||
}
|
||||
subgraphs = {
|
||||
k: v
|
||||
for k, v in zip(
|
||||
subpregels,
|
||||
await asyncio.gather(
|
||||
*(
|
||||
p.aget_graph(
|
||||
config,
|
||||
xray=xray
|
||||
if isinstance(xray, bool) or xray <= 0
|
||||
else xray - 1,
|
||||
)
|
||||
for p in subpregels.values()
|
||||
)
|
||||
),
|
||||
)
|
||||
}
|
||||
else:
|
||||
subgraphs = {}
|
||||
|
||||
return draw_graph(
|
||||
merge_configs(self.config, config),
|
||||
nodes=self.nodes,
|
||||
specs=self.channels,
|
||||
input_channels=self.input_channels,
|
||||
interrupt_after_nodes=self.interrupt_after_nodes,
|
||||
interrupt_before_nodes=self.interrupt_before_nodes,
|
||||
trigger_to_nodes=self.trigger_to_nodes,
|
||||
checkpointer=self.checkpointer,
|
||||
subgraphs=subgraphs,
|
||||
)
|
||||
|
||||
def _repr_mimebundle_(self, **kwargs: Any) -> dict[str, Any]:
|
||||
"""Mime bundle used by Jupyter to display the graph"""
|
||||
return {
|
||||
"text/plain": repr(self),
|
||||
"image/png": self.get_graph().draw_mermaid_png(),
|
||||
}
|
||||
|
||||
def copy(self, update: dict[str, Any] | None = None) -> Self:
|
||||
attrs = {**self.__dict__, **(update or {})}
|
||||
|
||||
@@ -0,0 +1,211 @@
|
||||
from collections import defaultdict
|
||||
from typing import Any, Mapping, Optional, Sequence, Union, cast
|
||||
|
||||
from langchain_core.runnables.config import RunnableConfig
|
||||
from langchain_core.runnables.graph import Graph, Node
|
||||
|
||||
from langgraph.channels.base import BaseChannel
|
||||
from langgraph.checkpoint.base import BaseCheckpointSaver
|
||||
from langgraph.constants import CONF, CONFIG_KEY_SEND, END, INPUT, START
|
||||
from langgraph.managed.base import ManagedValueSpec
|
||||
from langgraph.pregel.algo import (
|
||||
PregelTaskWrites,
|
||||
apply_writes,
|
||||
increment,
|
||||
prepare_next_tasks,
|
||||
)
|
||||
from langgraph.pregel.checkpoint import empty_checkpoint
|
||||
from langgraph.pregel.io import map_input
|
||||
from langgraph.pregel.manager import ChannelsManager
|
||||
from langgraph.pregel.read import PregelNode
|
||||
from langgraph.pregel.write import ChannelWrite
|
||||
from langgraph.types import All, Checkpointer, LoopProtocol
|
||||
|
||||
|
||||
def draw_graph(
|
||||
config: RunnableConfig,
|
||||
*,
|
||||
nodes: dict[str, PregelNode],
|
||||
specs: dict[str, Union[BaseChannel, ManagedValueSpec]],
|
||||
input_channels: Union[str, Sequence[str]],
|
||||
interrupt_after_nodes: Union[All, Sequence[str]],
|
||||
interrupt_before_nodes: Union[All, Sequence[str]],
|
||||
trigger_to_nodes: Optional[Mapping[str, Sequence[str]]],
|
||||
checkpointer: Checkpointer,
|
||||
subgraphs: dict[str, Graph],
|
||||
) -> Graph:
|
||||
"""Get the graph for this Pregel instance.
|
||||
|
||||
Args:
|
||||
config: The configuration to use for the graph.
|
||||
subgraphs: The subgraphs to include in the graph.
|
||||
checkpointer: The checkpointer to use for the graph.
|
||||
|
||||
Returns:
|
||||
The graph for this Pregel instance.
|
||||
"""
|
||||
# (src, dest, is_conditional, label)
|
||||
edges: set[tuple[str, str, bool, Optional[str]]] = set()
|
||||
|
||||
step = -1
|
||||
checkpoint = empty_checkpoint()
|
||||
get_next_version = (
|
||||
checkpointer.get_next_version
|
||||
if isinstance(checkpointer, BaseCheckpointSaver)
|
||||
else increment
|
||||
)
|
||||
with ChannelsManager(
|
||||
specs,
|
||||
checkpoint,
|
||||
LoopProtocol(step=step, stop=-1, config=config),
|
||||
skip_context=True,
|
||||
) as (channels, managed):
|
||||
static_seen: set[Any] = set()
|
||||
sources: dict[str, set[tuple[str, bool, Optional[str]]]] = {}
|
||||
step_sources: dict[str, set[tuple[str, bool, Optional[str]]]] = {}
|
||||
# remove node mappers
|
||||
nodes = {
|
||||
k: v.copy(update={"mapper": None}) if v.mapper is not None else v
|
||||
for k, v in nodes.items()
|
||||
}
|
||||
# apply input writes
|
||||
input_writes = list(map_input(input_channels, {}))
|
||||
_, updated_channels = apply_writes(
|
||||
checkpoint,
|
||||
channels,
|
||||
[
|
||||
PregelTaskWrites((), INPUT, input_writes, []),
|
||||
],
|
||||
get_next_version,
|
||||
)
|
||||
# prepare first tasks
|
||||
tasks = prepare_next_tasks(
|
||||
checkpoint,
|
||||
[],
|
||||
nodes,
|
||||
channels,
|
||||
managed,
|
||||
config,
|
||||
step,
|
||||
for_execution=True,
|
||||
store=None,
|
||||
checkpointer=None,
|
||||
manager=None,
|
||||
trigger_to_nodes=trigger_to_nodes,
|
||||
updated_channels=updated_channels,
|
||||
)
|
||||
start_tasks = tasks
|
||||
# run the pregel loop
|
||||
while tasks:
|
||||
conditionals: dict[tuple[str, str, Any], Optional[str]] = {}
|
||||
# run task writers
|
||||
for task in tasks.values():
|
||||
for w in task.writers:
|
||||
# apply regular writes
|
||||
if isinstance(w, ChannelWrite):
|
||||
w.invoke(None, task.config)
|
||||
# apply conditional writes declared for static analysis, only once
|
||||
if w not in static_seen:
|
||||
static_seen.add(w)
|
||||
# apply static writes
|
||||
if writes := ChannelWrite.get_static_writes(w):
|
||||
conditionals.update(
|
||||
{(task.name, *t[:2]): t[2] for t in writes}
|
||||
)
|
||||
task.config[CONF][CONFIG_KEY_SEND]([t[:2] for t in writes])
|
||||
# collect sources
|
||||
step_sources = {
|
||||
task.name: {
|
||||
(
|
||||
w[0],
|
||||
(task.name, *w) in conditionals,
|
||||
conditionals.get((task.name, *w)),
|
||||
)
|
||||
for w in task.writes
|
||||
}
|
||||
for task in tasks.values()
|
||||
}
|
||||
sources.update(step_sources)
|
||||
# invert triggers
|
||||
trigger_to_sources: dict[str, set[tuple[str, bool, Optional[str]]]] = (
|
||||
defaultdict(set)
|
||||
)
|
||||
for src, triggers in sources.items():
|
||||
for trigger, cond, label in triggers:
|
||||
trigger_to_sources[trigger].add((src, cond, label))
|
||||
# apply writes
|
||||
_, updated_channels = apply_writes(
|
||||
checkpoint, channels, tasks.values(), get_next_version
|
||||
)
|
||||
# prepare next tasks
|
||||
tasks = prepare_next_tasks(
|
||||
checkpoint,
|
||||
[],
|
||||
nodes,
|
||||
channels,
|
||||
managed,
|
||||
config,
|
||||
step,
|
||||
for_execution=True,
|
||||
store=None,
|
||||
checkpointer=None,
|
||||
manager=None,
|
||||
trigger_to_nodes=trigger_to_nodes,
|
||||
updated_channels=updated_channels,
|
||||
)
|
||||
# collect edges
|
||||
for task in tasks.values():
|
||||
for trigger in task.triggers:
|
||||
for src, cond, label in sorted(trigger_to_sources[trigger]):
|
||||
edges.add((src, task.name, cond, label))
|
||||
# assemble the graph
|
||||
graph = Graph()
|
||||
# add nodes
|
||||
for name, node in nodes.items():
|
||||
metadata = dict(node.metadata or {})
|
||||
if name in interrupt_before_nodes and name in interrupt_after_nodes:
|
||||
metadata["__interrupt"] = "before,after"
|
||||
elif name in interrupt_before_nodes:
|
||||
metadata["__interrupt"] = "before"
|
||||
elif name in interrupt_after_nodes:
|
||||
metadata["__interrupt"] = "after"
|
||||
graph.add_node(node.bound, name, metadata=metadata or None)
|
||||
# add start node
|
||||
if START not in nodes:
|
||||
graph.add_node(None, START)
|
||||
for task in start_tasks.values():
|
||||
graph.add_edge(graph.nodes[START], graph.nodes[task.name])
|
||||
# add discovered edges
|
||||
for src, dest, is_conditional, label in sorted(edges):
|
||||
graph.add_edge(
|
||||
graph.nodes[src],
|
||||
graph.nodes[dest],
|
||||
data=label if label != dest else None,
|
||||
conditional=is_conditional,
|
||||
)
|
||||
# add end edges
|
||||
if step_sources:
|
||||
end = graph.add_node(None, END)
|
||||
termini = {d for _, d, _, _ in edges}.difference(s for s, _, _, _ in edges)
|
||||
for src in sorted(termini.union(step_sources)):
|
||||
graph.add_edge(graph.nodes[src], end, conditional=src not in termini)
|
||||
# replace subgraphs
|
||||
for name, subgraph in subgraphs.items():
|
||||
subgraph.trim_first_node()
|
||||
subgraph.trim_last_node()
|
||||
if (
|
||||
len(subgraph.nodes) > 1
|
||||
and name in graph.nodes
|
||||
and subgraph.first_node()
|
||||
and subgraph.last_node()
|
||||
):
|
||||
# replace the node with the subgraph
|
||||
graph.nodes.pop(name)
|
||||
first, last = graph.extend(subgraph, prefix=name)
|
||||
for idx, edge in enumerate(graph.edges):
|
||||
if edge.source == name:
|
||||
graph.edges[idx] = edge.copy(source=cast(Node, last).id)
|
||||
elif edge.target == name:
|
||||
graph.edges[idx] = edge.copy(target=cast(Node, first).id)
|
||||
|
||||
return graph
|
||||
@@ -13,7 +13,7 @@ from typing import (
|
||||
from langchain_core.runnables import Runnable, RunnableConfig
|
||||
from langchain_core.runnables.utils import ConfigurableFieldSpec
|
||||
|
||||
from langgraph.constants import CONF, CONFIG_KEY_SEND, TASKS, Send
|
||||
from langgraph.constants import CONF, CONFIG_KEY_SEND, MISSING, TASKS, Send
|
||||
from langgraph.errors import InvalidUpdateError
|
||||
from langgraph.utils.runnable import RunnableCallable
|
||||
|
||||
@@ -40,6 +40,8 @@ class ChannelWriteTupleEntry(NamedTuple):
|
||||
"""Function to extract tuples from value."""
|
||||
value: Any = PASSTHROUGH
|
||||
"""Value to write, or PASSTHROUGH to use the input."""
|
||||
static: Optional[Sequence[tuple[str, Any, Optional[str]]]] = None
|
||||
"""Optional, declared writes for static analysis."""
|
||||
|
||||
|
||||
class ChannelWrite(RunnableCallable):
|
||||
@@ -118,6 +120,7 @@ class ChannelWrite(RunnableCallable):
|
||||
def do_write(
|
||||
config: RunnableConfig,
|
||||
writes: Sequence[ChannelWriteEntry | ChannelWriteTupleEntry | Send],
|
||||
allow_passthrough: bool = True,
|
||||
require_at_least_one_of: Sequence[str] | None = None, # ignored
|
||||
) -> None:
|
||||
# validate
|
||||
@@ -127,46 +130,80 @@ class ChannelWrite(RunnableCallable):
|
||||
raise InvalidUpdateError(
|
||||
"Cannot write to the reserved channel TASKS"
|
||||
)
|
||||
if w.value is PASSTHROUGH:
|
||||
if w.value is PASSTHROUGH and not allow_passthrough:
|
||||
raise InvalidUpdateError("PASSTHROUGH value must be replaced")
|
||||
if isinstance(w, ChannelWriteTupleEntry):
|
||||
if w.value is PASSTHROUGH:
|
||||
if w.value is PASSTHROUGH and not allow_passthrough:
|
||||
raise InvalidUpdateError("PASSTHROUGH value must be replaced")
|
||||
# assemble writes
|
||||
tuples: list[tuple[str, Any]] = []
|
||||
for w in writes:
|
||||
if isinstance(w, Send):
|
||||
tuples.append((TASKS, w))
|
||||
elif isinstance(w, ChannelWriteTupleEntry):
|
||||
if ww := w.mapper(w.value):
|
||||
tuples.extend(ww)
|
||||
elif isinstance(w, ChannelWriteEntry):
|
||||
value = w.mapper(w.value) if w.mapper is not None else w.value
|
||||
if value is SKIP_WRITE:
|
||||
continue
|
||||
if w.skip_none and value is None:
|
||||
continue
|
||||
tuples.append((w.channel, value))
|
||||
else:
|
||||
raise ValueError(f"Invalid write entry: {w}")
|
||||
# if we want to persist writes found before hitting a ParentCommand
|
||||
# can move this to a finally block
|
||||
write: TYPE_SEND = config[CONF][CONFIG_KEY_SEND]
|
||||
write(tuples)
|
||||
write(_assemble_writes(writes))
|
||||
|
||||
@staticmethod
|
||||
def is_writer(runnable: Runnable) -> bool:
|
||||
"""Used by PregelNode to distinguish between writers and other runnables."""
|
||||
return (
|
||||
isinstance(runnable, ChannelWrite)
|
||||
or getattr(runnable, "_is_channel_writer", False) is True
|
||||
or getattr(runnable, "_is_channel_writer", MISSING) is not MISSING
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def register_writer(runnable: R) -> R:
|
||||
def get_static_writes(
|
||||
runnable: Runnable,
|
||||
) -> Optional[Sequence[tuple[str, Any, Optional[str]]]]:
|
||||
"""Used to get conditional writes a writer declares for static analysis."""
|
||||
if isinstance(runnable, ChannelWrite):
|
||||
return [
|
||||
w
|
||||
for entry in runnable.writes
|
||||
if isinstance(entry, ChannelWriteTupleEntry) and entry.static
|
||||
for w in entry.static
|
||||
] or None
|
||||
elif writes := getattr(runnable, "_is_channel_writer", MISSING):
|
||||
if writes is not MISSING:
|
||||
writes = cast(
|
||||
Sequence[tuple[Union[ChannelWriteEntry, Send], Optional[str]]],
|
||||
writes,
|
||||
)
|
||||
entries = [e for e, _ in writes]
|
||||
labels = [la for _, la in writes]
|
||||
return [(*t, la) for t, la in zip(_assemble_writes(entries), labels)]
|
||||
|
||||
@staticmethod
|
||||
def register_writer(
|
||||
runnable: R,
|
||||
static: Optional[
|
||||
Sequence[tuple[Union[ChannelWriteEntry, Send], Optional[str]]]
|
||||
] = None,
|
||||
) -> R:
|
||||
"""Used to mark a runnable as a writer, so that it can be detected by is_writer.
|
||||
Instances of ChannelWrite are automatically marked as writers."""
|
||||
Instances of ChannelWrite are automatically marked as writers.
|
||||
Optionally, a list of declared writes can be passed for static analysis."""
|
||||
# using object.__setattr__ to work around objects that override __setattr__
|
||||
# eg. pydantic models and dataclasses
|
||||
object.__setattr__(runnable, "_is_channel_writer", True)
|
||||
object.__setattr__(runnable, "_is_channel_writer", static)
|
||||
return runnable
|
||||
|
||||
|
||||
def _assemble_writes(
|
||||
writes: Sequence[Union[ChannelWriteEntry, ChannelWriteTupleEntry, Send]],
|
||||
) -> list[tuple[str, Any]]:
|
||||
"""Assembles the writes into a list of tuples."""
|
||||
tuples: list[tuple[str, Any]] = []
|
||||
for w in writes:
|
||||
if isinstance(w, Send):
|
||||
tuples.append((TASKS, w))
|
||||
elif isinstance(w, ChannelWriteTupleEntry):
|
||||
if ww := w.mapper(w.value):
|
||||
tuples.extend(ww)
|
||||
elif isinstance(w, ChannelWriteEntry):
|
||||
value = w.mapper(w.value) if w.mapper is not None else w.value
|
||||
if value is SKIP_WRITE:
|
||||
continue
|
||||
if w.skip_none and value is None:
|
||||
continue
|
||||
tuples.append((w.channel, value))
|
||||
else:
|
||||
raise ValueError(f"Invalid write entry: {w}")
|
||||
return tuples
|
||||
|
||||
Generated
+4
-4
@@ -1,4 +1,4 @@
|
||||
# This file is automatically @generated by Poetry 2.0.1 and should not be changed by hand.
|
||||
# This file is automatically @generated by Poetry 2.0.0 and should not be changed by hand.
|
||||
|
||||
[[package]]
|
||||
name = "aiosqlite"
|
||||
@@ -1324,14 +1324,14 @@ files = [
|
||||
|
||||
[[package]]
|
||||
name = "langchain-core"
|
||||
version = "0.3.46"
|
||||
version = "0.3.55"
|
||||
description = "Building applications with LLMs through composability"
|
||||
optional = false
|
||||
python-versions = "<4.0,>=3.9"
|
||||
groups = ["main", "dev"]
|
||||
files = [
|
||||
{file = "langchain_core-0.3.46-py3-none-any.whl", hash = "sha256:28b5689fc347975ea520b5364ab4aee5567e661553bbee5e97cabf4596c28ce0"},
|
||||
{file = "langchain_core-0.3.46.tar.gz", hash = "sha256:5fca010eeb0a427be5aa8a8525e2112995dde790c584cef165be7c5e0ee1c2b5"},
|
||||
{file = "langchain_core-0.3.55-py3-none-any.whl", hash = "sha256:b3cb36bf37755a616158a79866657c6697b43a2f7c69dd723ce425f1c76c1baa"},
|
||||
{file = "langchain_core-0.3.55.tar.gz", hash = "sha256:0f2b3e311621116a83510c70b0ac9d959030a0a457a69483535cff18501fedc9"},
|
||||
]
|
||||
|
||||
[package.dependencies]
|
||||
|
||||
File diff suppressed because one or more lines are too long
@@ -1,151 +0,0 @@
|
||||
# serializer version: 1
|
||||
# name: test_weather_subgraph[memory]
|
||||
'''
|
||||
%%{init: {'flowchart': {'curve': 'linear'}}}%%
|
||||
graph TD;
|
||||
__start__([<p>__start__</p>]):::first
|
||||
router_node(router_node)
|
||||
normal_llm_node(normal_llm_node)
|
||||
__end__([<p>__end__</p>]):::last
|
||||
__start__ --> router_node;
|
||||
normal_llm_node --> __end__;
|
||||
weather_graph_weather_node --> __end__;
|
||||
router_node -.-> normal_llm_node;
|
||||
router_node -.-> weather_graph_model_node;
|
||||
router_node -.-> __end__;
|
||||
subgraph weather_graph
|
||||
weather_graph_model_node(model_node)
|
||||
weather_graph_weather_node(weather_node<hr/><small><em>__interrupt = before</em></small>)
|
||||
weather_graph_model_node --> weather_graph_weather_node;
|
||||
end
|
||||
classDef default fill:#f2f0ff,line-height:1.2
|
||||
classDef first fill-opacity:0
|
||||
classDef last fill:#bfb6fc
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_weather_subgraph[postgres_aio]
|
||||
'''
|
||||
%%{init: {'flowchart': {'curve': 'linear'}}}%%
|
||||
graph TD;
|
||||
__start__([<p>__start__</p>]):::first
|
||||
router_node(router_node)
|
||||
normal_llm_node(normal_llm_node)
|
||||
__end__([<p>__end__</p>]):::last
|
||||
__start__ --> router_node;
|
||||
normal_llm_node --> __end__;
|
||||
weather_graph_weather_node --> __end__;
|
||||
router_node -.-> normal_llm_node;
|
||||
router_node -.-> weather_graph_model_node;
|
||||
router_node -.-> __end__;
|
||||
subgraph weather_graph
|
||||
weather_graph_model_node(model_node)
|
||||
weather_graph_weather_node(weather_node<hr/><small><em>__interrupt = before</em></small>)
|
||||
weather_graph_model_node --> weather_graph_weather_node;
|
||||
end
|
||||
classDef default fill:#f2f0ff,line-height:1.2
|
||||
classDef first fill-opacity:0
|
||||
classDef last fill:#bfb6fc
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_weather_subgraph[postgres_aio_pipe]
|
||||
'''
|
||||
%%{init: {'flowchart': {'curve': 'linear'}}}%%
|
||||
graph TD;
|
||||
__start__([<p>__start__</p>]):::first
|
||||
router_node(router_node)
|
||||
normal_llm_node(normal_llm_node)
|
||||
__end__([<p>__end__</p>]):::last
|
||||
__start__ --> router_node;
|
||||
normal_llm_node --> __end__;
|
||||
weather_graph_weather_node --> __end__;
|
||||
router_node -.-> normal_llm_node;
|
||||
router_node -.-> weather_graph_model_node;
|
||||
router_node -.-> __end__;
|
||||
subgraph weather_graph
|
||||
weather_graph_model_node(model_node)
|
||||
weather_graph_weather_node(weather_node<hr/><small><em>__interrupt = before</em></small>)
|
||||
weather_graph_model_node --> weather_graph_weather_node;
|
||||
end
|
||||
classDef default fill:#f2f0ff,line-height:1.2
|
||||
classDef first fill-opacity:0
|
||||
classDef last fill:#bfb6fc
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_weather_subgraph[postgres_aio_pool]
|
||||
'''
|
||||
%%{init: {'flowchart': {'curve': 'linear'}}}%%
|
||||
graph TD;
|
||||
__start__([<p>__start__</p>]):::first
|
||||
router_node(router_node)
|
||||
normal_llm_node(normal_llm_node)
|
||||
__end__([<p>__end__</p>]):::last
|
||||
__start__ --> router_node;
|
||||
normal_llm_node --> __end__;
|
||||
weather_graph_weather_node --> __end__;
|
||||
router_node -.-> normal_llm_node;
|
||||
router_node -.-> weather_graph_model_node;
|
||||
router_node -.-> __end__;
|
||||
subgraph weather_graph
|
||||
weather_graph_model_node(model_node)
|
||||
weather_graph_weather_node(weather_node<hr/><small><em>__interrupt = before</em></small>)
|
||||
weather_graph_model_node --> weather_graph_weather_node;
|
||||
end
|
||||
classDef default fill:#f2f0ff,line-height:1.2
|
||||
classDef first fill-opacity:0
|
||||
classDef last fill:#bfb6fc
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_weather_subgraph[postgres_aio_shallow]
|
||||
'''
|
||||
%%{init: {'flowchart': {'curve': 'linear'}}}%%
|
||||
graph TD;
|
||||
__start__([<p>__start__</p>]):::first
|
||||
router_node(router_node)
|
||||
normal_llm_node(normal_llm_node)
|
||||
__end__([<p>__end__</p>]):::last
|
||||
__start__ --> router_node;
|
||||
normal_llm_node --> __end__;
|
||||
weather_graph_weather_node --> __end__;
|
||||
router_node -.-> normal_llm_node;
|
||||
router_node -.-> weather_graph_model_node;
|
||||
router_node -.-> __end__;
|
||||
subgraph weather_graph
|
||||
weather_graph_model_node(model_node)
|
||||
weather_graph_weather_node(weather_node<hr/><small><em>__interrupt = before</em></small>)
|
||||
weather_graph_model_node --> weather_graph_weather_node;
|
||||
end
|
||||
classDef default fill:#f2f0ff,line-height:1.2
|
||||
classDef first fill-opacity:0
|
||||
classDef last fill:#bfb6fc
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_weather_subgraph[sqlite_aio]
|
||||
'''
|
||||
%%{init: {'flowchart': {'curve': 'linear'}}}%%
|
||||
graph TD;
|
||||
__start__([<p>__start__</p>]):::first
|
||||
router_node(router_node)
|
||||
normal_llm_node(normal_llm_node)
|
||||
__end__([<p>__end__</p>]):::last
|
||||
__start__ --> router_node;
|
||||
normal_llm_node --> __end__;
|
||||
weather_graph_weather_node --> __end__;
|
||||
router_node -.-> normal_llm_node;
|
||||
router_node -.-> weather_graph_model_node;
|
||||
router_node -.-> __end__;
|
||||
subgraph weather_graph
|
||||
weather_graph_model_node(model_node)
|
||||
weather_graph_weather_node(weather_node<hr/><small><em>__interrupt = before</em></small>)
|
||||
weather_graph_model_node --> weather_graph_weather_node;
|
||||
end
|
||||
classDef default fill:#f2f0ff,line-height:1.2
|
||||
classDef first fill-opacity:0
|
||||
classDef last fill:#bfb6fc
|
||||
|
||||
'''
|
||||
# ---
|
||||
File diff suppressed because it is too large
Load Diff
@@ -4,11 +4,11 @@
|
||||
graph TD;
|
||||
__start__ --> rewrite_query;
|
||||
analyzer_one --> retriever_one;
|
||||
qa --> __end__;
|
||||
retriever_one --> qa;
|
||||
retriever_two --> qa;
|
||||
rewrite_query --> analyzer_one;
|
||||
rewrite_query -.-> retriever_two;
|
||||
qa --> __end__;
|
||||
|
||||
'''
|
||||
# ---
|
||||
@@ -120,695 +120,21 @@
|
||||
'type': 'object',
|
||||
})
|
||||
# ---
|
||||
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio]
|
||||
'''
|
||||
graph TD;
|
||||
__start__ --> rewrite_query;
|
||||
analyzer_one --> retriever_one;
|
||||
qa --> __end__;
|
||||
retriever_one --> qa;
|
||||
retriever_two --> qa;
|
||||
rewrite_query --> analyzer_one;
|
||||
rewrite_query -.-> retriever_two;
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio].1
|
||||
dict({
|
||||
'$defs': dict({
|
||||
'InnerObject': dict({
|
||||
'properties': dict({
|
||||
'yo': dict({
|
||||
'title': 'Yo',
|
||||
'type': 'integer',
|
||||
}),
|
||||
}),
|
||||
'required': list([
|
||||
'yo',
|
||||
]),
|
||||
'title': 'InnerObject',
|
||||
'type': 'object',
|
||||
}),
|
||||
}),
|
||||
'properties': dict({
|
||||
'answer': dict({
|
||||
'anyOf': list([
|
||||
dict({
|
||||
'type': 'string',
|
||||
}),
|
||||
dict({
|
||||
'type': 'null',
|
||||
}),
|
||||
]),
|
||||
'default': None,
|
||||
'title': 'Answer',
|
||||
}),
|
||||
'docs': dict({
|
||||
'items': dict({
|
||||
'type': 'string',
|
||||
}),
|
||||
'title': 'Docs',
|
||||
'type': 'array',
|
||||
}),
|
||||
'inner': dict({
|
||||
'$ref': '#/$defs/InnerObject',
|
||||
}),
|
||||
'query': dict({
|
||||
'title': 'Query',
|
||||
'type': 'string',
|
||||
}),
|
||||
}),
|
||||
'required': list([
|
||||
'query',
|
||||
'inner',
|
||||
'docs',
|
||||
]),
|
||||
'title': 'State',
|
||||
'type': 'object',
|
||||
})
|
||||
# ---
|
||||
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio].2
|
||||
dict({
|
||||
'$defs': dict({
|
||||
'InnerObject': dict({
|
||||
'properties': dict({
|
||||
'yo': dict({
|
||||
'title': 'Yo',
|
||||
'type': 'integer',
|
||||
}),
|
||||
}),
|
||||
'required': list([
|
||||
'yo',
|
||||
]),
|
||||
'title': 'InnerObject',
|
||||
'type': 'object',
|
||||
}),
|
||||
}),
|
||||
'properties': dict({
|
||||
'answer': dict({
|
||||
'anyOf': list([
|
||||
dict({
|
||||
'type': 'string',
|
||||
}),
|
||||
dict({
|
||||
'type': 'null',
|
||||
}),
|
||||
]),
|
||||
'default': None,
|
||||
'title': 'Answer',
|
||||
}),
|
||||
'docs': dict({
|
||||
'items': dict({
|
||||
'type': 'string',
|
||||
}),
|
||||
'title': 'Docs',
|
||||
'type': 'array',
|
||||
}),
|
||||
'inner': dict({
|
||||
'$ref': '#/$defs/InnerObject',
|
||||
}),
|
||||
'query': dict({
|
||||
'title': 'Query',
|
||||
'type': 'string',
|
||||
}),
|
||||
}),
|
||||
'required': list([
|
||||
'query',
|
||||
'inner',
|
||||
'docs',
|
||||
]),
|
||||
'title': 'State',
|
||||
'type': 'object',
|
||||
})
|
||||
# ---
|
||||
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio_pipe]
|
||||
'''
|
||||
graph TD;
|
||||
__start__ --> rewrite_query;
|
||||
analyzer_one --> retriever_one;
|
||||
qa --> __end__;
|
||||
retriever_one --> qa;
|
||||
retriever_two --> qa;
|
||||
rewrite_query --> analyzer_one;
|
||||
rewrite_query -.-> retriever_two;
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio_pipe].1
|
||||
dict({
|
||||
'$defs': dict({
|
||||
'InnerObject': dict({
|
||||
'properties': dict({
|
||||
'yo': dict({
|
||||
'title': 'Yo',
|
||||
'type': 'integer',
|
||||
}),
|
||||
}),
|
||||
'required': list([
|
||||
'yo',
|
||||
]),
|
||||
'title': 'InnerObject',
|
||||
'type': 'object',
|
||||
}),
|
||||
}),
|
||||
'properties': dict({
|
||||
'answer': dict({
|
||||
'anyOf': list([
|
||||
dict({
|
||||
'type': 'string',
|
||||
}),
|
||||
dict({
|
||||
'type': 'null',
|
||||
}),
|
||||
]),
|
||||
'default': None,
|
||||
'title': 'Answer',
|
||||
}),
|
||||
'docs': dict({
|
||||
'items': dict({
|
||||
'type': 'string',
|
||||
}),
|
||||
'title': 'Docs',
|
||||
'type': 'array',
|
||||
}),
|
||||
'inner': dict({
|
||||
'$ref': '#/$defs/InnerObject',
|
||||
}),
|
||||
'query': dict({
|
||||
'title': 'Query',
|
||||
'type': 'string',
|
||||
}),
|
||||
}),
|
||||
'required': list([
|
||||
'query',
|
||||
'inner',
|
||||
'docs',
|
||||
]),
|
||||
'title': 'State',
|
||||
'type': 'object',
|
||||
})
|
||||
# ---
|
||||
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio_pipe].2
|
||||
dict({
|
||||
'$defs': dict({
|
||||
'InnerObject': dict({
|
||||
'properties': dict({
|
||||
'yo': dict({
|
||||
'title': 'Yo',
|
||||
'type': 'integer',
|
||||
}),
|
||||
}),
|
||||
'required': list([
|
||||
'yo',
|
||||
]),
|
||||
'title': 'InnerObject',
|
||||
'type': 'object',
|
||||
}),
|
||||
}),
|
||||
'properties': dict({
|
||||
'answer': dict({
|
||||
'anyOf': list([
|
||||
dict({
|
||||
'type': 'string',
|
||||
}),
|
||||
dict({
|
||||
'type': 'null',
|
||||
}),
|
||||
]),
|
||||
'default': None,
|
||||
'title': 'Answer',
|
||||
}),
|
||||
'docs': dict({
|
||||
'items': dict({
|
||||
'type': 'string',
|
||||
}),
|
||||
'title': 'Docs',
|
||||
'type': 'array',
|
||||
}),
|
||||
'inner': dict({
|
||||
'$ref': '#/$defs/InnerObject',
|
||||
}),
|
||||
'query': dict({
|
||||
'title': 'Query',
|
||||
'type': 'string',
|
||||
}),
|
||||
}),
|
||||
'required': list([
|
||||
'query',
|
||||
'inner',
|
||||
'docs',
|
||||
]),
|
||||
'title': 'State',
|
||||
'type': 'object',
|
||||
})
|
||||
# ---
|
||||
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio_pool]
|
||||
'''
|
||||
graph TD;
|
||||
__start__ --> rewrite_query;
|
||||
analyzer_one --> retriever_one;
|
||||
qa --> __end__;
|
||||
retriever_one --> qa;
|
||||
retriever_two --> qa;
|
||||
rewrite_query --> analyzer_one;
|
||||
rewrite_query -.-> retriever_two;
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio_pool].1
|
||||
dict({
|
||||
'$defs': dict({
|
||||
'InnerObject': dict({
|
||||
'properties': dict({
|
||||
'yo': dict({
|
||||
'title': 'Yo',
|
||||
'type': 'integer',
|
||||
}),
|
||||
}),
|
||||
'required': list([
|
||||
'yo',
|
||||
]),
|
||||
'title': 'InnerObject',
|
||||
'type': 'object',
|
||||
}),
|
||||
}),
|
||||
'properties': dict({
|
||||
'answer': dict({
|
||||
'anyOf': list([
|
||||
dict({
|
||||
'type': 'string',
|
||||
}),
|
||||
dict({
|
||||
'type': 'null',
|
||||
}),
|
||||
]),
|
||||
'default': None,
|
||||
'title': 'Answer',
|
||||
}),
|
||||
'docs': dict({
|
||||
'items': dict({
|
||||
'type': 'string',
|
||||
}),
|
||||
'title': 'Docs',
|
||||
'type': 'array',
|
||||
}),
|
||||
'inner': dict({
|
||||
'$ref': '#/$defs/InnerObject',
|
||||
}),
|
||||
'query': dict({
|
||||
'title': 'Query',
|
||||
'type': 'string',
|
||||
}),
|
||||
}),
|
||||
'required': list([
|
||||
'query',
|
||||
'inner',
|
||||
'docs',
|
||||
]),
|
||||
'title': 'State',
|
||||
'type': 'object',
|
||||
})
|
||||
# ---
|
||||
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio_pool].2
|
||||
dict({
|
||||
'$defs': dict({
|
||||
'InnerObject': dict({
|
||||
'properties': dict({
|
||||
'yo': dict({
|
||||
'title': 'Yo',
|
||||
'type': 'integer',
|
||||
}),
|
||||
}),
|
||||
'required': list([
|
||||
'yo',
|
||||
]),
|
||||
'title': 'InnerObject',
|
||||
'type': 'object',
|
||||
}),
|
||||
}),
|
||||
'properties': dict({
|
||||
'answer': dict({
|
||||
'anyOf': list([
|
||||
dict({
|
||||
'type': 'string',
|
||||
}),
|
||||
dict({
|
||||
'type': 'null',
|
||||
}),
|
||||
]),
|
||||
'default': None,
|
||||
'title': 'Answer',
|
||||
}),
|
||||
'docs': dict({
|
||||
'items': dict({
|
||||
'type': 'string',
|
||||
}),
|
||||
'title': 'Docs',
|
||||
'type': 'array',
|
||||
}),
|
||||
'inner': dict({
|
||||
'$ref': '#/$defs/InnerObject',
|
||||
}),
|
||||
'query': dict({
|
||||
'title': 'Query',
|
||||
'type': 'string',
|
||||
}),
|
||||
}),
|
||||
'required': list([
|
||||
'query',
|
||||
'inner',
|
||||
'docs',
|
||||
]),
|
||||
'title': 'State',
|
||||
'type': 'object',
|
||||
})
|
||||
# ---
|
||||
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio_shallow]
|
||||
'''
|
||||
graph TD;
|
||||
__start__ --> rewrite_query;
|
||||
analyzer_one --> retriever_one;
|
||||
qa --> __end__;
|
||||
retriever_one --> qa;
|
||||
retriever_two --> qa;
|
||||
rewrite_query --> analyzer_one;
|
||||
rewrite_query -.-> retriever_two;
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio_shallow].1
|
||||
dict({
|
||||
'$defs': dict({
|
||||
'InnerObject': dict({
|
||||
'properties': dict({
|
||||
'yo': dict({
|
||||
'title': 'Yo',
|
||||
'type': 'integer',
|
||||
}),
|
||||
}),
|
||||
'required': list([
|
||||
'yo',
|
||||
]),
|
||||
'title': 'InnerObject',
|
||||
'type': 'object',
|
||||
}),
|
||||
}),
|
||||
'properties': dict({
|
||||
'answer': dict({
|
||||
'anyOf': list([
|
||||
dict({
|
||||
'type': 'string',
|
||||
}),
|
||||
dict({
|
||||
'type': 'null',
|
||||
}),
|
||||
]),
|
||||
'default': None,
|
||||
'title': 'Answer',
|
||||
}),
|
||||
'docs': dict({
|
||||
'items': dict({
|
||||
'type': 'string',
|
||||
}),
|
||||
'title': 'Docs',
|
||||
'type': 'array',
|
||||
}),
|
||||
'inner': dict({
|
||||
'$ref': '#/$defs/InnerObject',
|
||||
}),
|
||||
'query': dict({
|
||||
'title': 'Query',
|
||||
'type': 'string',
|
||||
}),
|
||||
}),
|
||||
'required': list([
|
||||
'query',
|
||||
'inner',
|
||||
'docs',
|
||||
]),
|
||||
'title': 'State',
|
||||
'type': 'object',
|
||||
})
|
||||
# ---
|
||||
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio_shallow].2
|
||||
dict({
|
||||
'$defs': dict({
|
||||
'InnerObject': dict({
|
||||
'properties': dict({
|
||||
'yo': dict({
|
||||
'title': 'Yo',
|
||||
'type': 'integer',
|
||||
}),
|
||||
}),
|
||||
'required': list([
|
||||
'yo',
|
||||
]),
|
||||
'title': 'InnerObject',
|
||||
'type': 'object',
|
||||
}),
|
||||
}),
|
||||
'properties': dict({
|
||||
'answer': dict({
|
||||
'anyOf': list([
|
||||
dict({
|
||||
'type': 'string',
|
||||
}),
|
||||
dict({
|
||||
'type': 'null',
|
||||
}),
|
||||
]),
|
||||
'default': None,
|
||||
'title': 'Answer',
|
||||
}),
|
||||
'docs': dict({
|
||||
'items': dict({
|
||||
'type': 'string',
|
||||
}),
|
||||
'title': 'Docs',
|
||||
'type': 'array',
|
||||
}),
|
||||
'inner': dict({
|
||||
'$ref': '#/$defs/InnerObject',
|
||||
}),
|
||||
'query': dict({
|
||||
'title': 'Query',
|
||||
'type': 'string',
|
||||
}),
|
||||
}),
|
||||
'required': list([
|
||||
'query',
|
||||
'inner',
|
||||
'docs',
|
||||
]),
|
||||
'title': 'State',
|
||||
'type': 'object',
|
||||
})
|
||||
# ---
|
||||
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[sqlite_aio]
|
||||
'''
|
||||
graph TD;
|
||||
__start__ --> rewrite_query;
|
||||
analyzer_one --> retriever_one;
|
||||
qa --> __end__;
|
||||
retriever_one --> qa;
|
||||
retriever_two --> qa;
|
||||
rewrite_query --> analyzer_one;
|
||||
rewrite_query -.-> retriever_two;
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[sqlite_aio].1
|
||||
dict({
|
||||
'$defs': dict({
|
||||
'InnerObject': dict({
|
||||
'properties': dict({
|
||||
'yo': dict({
|
||||
'title': 'Yo',
|
||||
'type': 'integer',
|
||||
}),
|
||||
}),
|
||||
'required': list([
|
||||
'yo',
|
||||
]),
|
||||
'title': 'InnerObject',
|
||||
'type': 'object',
|
||||
}),
|
||||
}),
|
||||
'properties': dict({
|
||||
'answer': dict({
|
||||
'anyOf': list([
|
||||
dict({
|
||||
'type': 'string',
|
||||
}),
|
||||
dict({
|
||||
'type': 'null',
|
||||
}),
|
||||
]),
|
||||
'default': None,
|
||||
'title': 'Answer',
|
||||
}),
|
||||
'docs': dict({
|
||||
'items': dict({
|
||||
'type': 'string',
|
||||
}),
|
||||
'title': 'Docs',
|
||||
'type': 'array',
|
||||
}),
|
||||
'inner': dict({
|
||||
'$ref': '#/$defs/InnerObject',
|
||||
}),
|
||||
'query': dict({
|
||||
'title': 'Query',
|
||||
'type': 'string',
|
||||
}),
|
||||
}),
|
||||
'required': list([
|
||||
'query',
|
||||
'inner',
|
||||
'docs',
|
||||
]),
|
||||
'title': 'State',
|
||||
'type': 'object',
|
||||
})
|
||||
# ---
|
||||
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[sqlite_aio].2
|
||||
dict({
|
||||
'$defs': dict({
|
||||
'InnerObject': dict({
|
||||
'properties': dict({
|
||||
'yo': dict({
|
||||
'title': 'Yo',
|
||||
'type': 'integer',
|
||||
}),
|
||||
}),
|
||||
'required': list([
|
||||
'yo',
|
||||
]),
|
||||
'title': 'InnerObject',
|
||||
'type': 'object',
|
||||
}),
|
||||
}),
|
||||
'properties': dict({
|
||||
'answer': dict({
|
||||
'anyOf': list([
|
||||
dict({
|
||||
'type': 'string',
|
||||
}),
|
||||
dict({
|
||||
'type': 'null',
|
||||
}),
|
||||
]),
|
||||
'default': None,
|
||||
'title': 'Answer',
|
||||
}),
|
||||
'docs': dict({
|
||||
'items': dict({
|
||||
'type': 'string',
|
||||
}),
|
||||
'title': 'Docs',
|
||||
'type': 'array',
|
||||
}),
|
||||
'inner': dict({
|
||||
'$ref': '#/$defs/InnerObject',
|
||||
}),
|
||||
'query': dict({
|
||||
'title': 'Query',
|
||||
'type': 'string',
|
||||
}),
|
||||
}),
|
||||
'required': list([
|
||||
'query',
|
||||
'inner',
|
||||
'docs',
|
||||
]),
|
||||
'title': 'State',
|
||||
'type': 'object',
|
||||
})
|
||||
# ---
|
||||
# name: test_send_react_interrupt_control[memory]
|
||||
'''
|
||||
%%{init: {'flowchart': {'curve': 'linear'}}}%%
|
||||
---
|
||||
config:
|
||||
flowchart:
|
||||
curve: linear
|
||||
---
|
||||
graph TD;
|
||||
__start__([<p>__start__</p>]):::first
|
||||
agent(agent)
|
||||
foo([foo]):::last
|
||||
__start__ --> agent;
|
||||
agent -.-> foo;
|
||||
classDef default fill:#f2f0ff,line-height:1.2
|
||||
classDef first fill-opacity:0
|
||||
classDef last fill:#bfb6fc
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_send_react_interrupt_control[postgres_aio]
|
||||
'''
|
||||
%%{init: {'flowchart': {'curve': 'linear'}}}%%
|
||||
graph TD;
|
||||
__start__([<p>__start__</p>]):::first
|
||||
agent(agent)
|
||||
foo([foo]):::last
|
||||
__start__ --> agent;
|
||||
agent -.-> foo;
|
||||
classDef default fill:#f2f0ff,line-height:1.2
|
||||
classDef first fill-opacity:0
|
||||
classDef last fill:#bfb6fc
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_send_react_interrupt_control[postgres_aio_pipe]
|
||||
'''
|
||||
%%{init: {'flowchart': {'curve': 'linear'}}}%%
|
||||
graph TD;
|
||||
__start__([<p>__start__</p>]):::first
|
||||
agent(agent)
|
||||
foo([foo]):::last
|
||||
__start__ --> agent;
|
||||
agent -.-> foo;
|
||||
classDef default fill:#f2f0ff,line-height:1.2
|
||||
classDef first fill-opacity:0
|
||||
classDef last fill:#bfb6fc
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_send_react_interrupt_control[postgres_aio_pool]
|
||||
'''
|
||||
%%{init: {'flowchart': {'curve': 'linear'}}}%%
|
||||
graph TD;
|
||||
__start__([<p>__start__</p>]):::first
|
||||
agent(agent)
|
||||
foo([foo]):::last
|
||||
__start__ --> agent;
|
||||
agent -.-> foo;
|
||||
classDef default fill:#f2f0ff,line-height:1.2
|
||||
classDef first fill-opacity:0
|
||||
classDef last fill:#bfb6fc
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_send_react_interrupt_control[postgres_aio_shallow]
|
||||
'''
|
||||
%%{init: {'flowchart': {'curve': 'linear'}}}%%
|
||||
graph TD;
|
||||
__start__([<p>__start__</p>]):::first
|
||||
agent(agent)
|
||||
foo([foo]):::last
|
||||
__start__ --> agent;
|
||||
agent -.-> foo;
|
||||
classDef default fill:#f2f0ff,line-height:1.2
|
||||
classDef first fill-opacity:0
|
||||
classDef last fill:#bfb6fc
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_send_react_interrupt_control[sqlite_aio]
|
||||
'''
|
||||
%%{init: {'flowchart': {'curve': 'linear'}}}%%
|
||||
graph TD;
|
||||
__start__([<p>__start__</p>]):::first
|
||||
agent(agent)
|
||||
foo([foo]):::last
|
||||
foo(foo)
|
||||
__end__([<p>__end__</p>]):::last
|
||||
__start__ --> agent;
|
||||
agent -.-> foo;
|
||||
foo --> __end__;
|
||||
classDef default fill:#f2f0ff,line-height:1.2
|
||||
classDef first fill-opacity:0
|
||||
classDef last fill:#bfb6fc
|
||||
|
||||
@@ -588,12 +588,10 @@ def test_conditional_graph(
|
||||
|
||||
app = workflow.compile()
|
||||
|
||||
if SHOULD_CHECK_SNAPSHOTS:
|
||||
if SHOULD_CHECK_SNAPSHOTS and checkpointer_name == "memory":
|
||||
assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot
|
||||
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
|
||||
assert app.get_graph().draw_mermaid() == snapshot
|
||||
assert json.dumps(app.get_graph(xray=True).to_json(), indent=2) == snapshot
|
||||
assert app.get_graph(xray=True).draw_mermaid(with_styles=False) == snapshot
|
||||
|
||||
assert app.invoke({"input": "what is weather in sf"}) == {
|
||||
"input": "what is weather in sf",
|
||||
@@ -723,10 +721,6 @@ def test_conditional_graph(
|
||||
)
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
|
||||
if SHOULD_CHECK_SNAPSHOTS:
|
||||
assert app_w_interrupt.get_graph().to_json() == snapshot
|
||||
assert app_w_interrupt.get_graph().draw_mermaid() == snapshot
|
||||
|
||||
assert [
|
||||
c for c in app_w_interrupt.stream({"input": "what is weather in sf"}, config)
|
||||
] == [
|
||||
@@ -1539,7 +1533,7 @@ def test_conditional_state_graph(
|
||||
|
||||
app = workflow.compile()
|
||||
|
||||
if SHOULD_CHECK_SNAPSHOTS:
|
||||
if SHOULD_CHECK_SNAPSHOTS and checkpointer_name == "memory":
|
||||
assert json.dumps(app.get_input_schema().model_json_schema()) == snapshot
|
||||
assert json.dumps(app.get_output_schema().model_json_schema()) == snapshot
|
||||
assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot
|
||||
@@ -3775,7 +3769,7 @@ def test_message_graph(
|
||||
# meaning you can use it as you would any other runnable
|
||||
app = workflow.compile()
|
||||
|
||||
if SHOULD_CHECK_SNAPSHOTS:
|
||||
if SHOULD_CHECK_SNAPSHOTS and checkpointer_name == "memory":
|
||||
assert json.dumps(app.get_input_schema().model_json_schema()) == snapshot
|
||||
assert json.dumps(app.get_output_schema().model_json_schema()) == snapshot
|
||||
assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot
|
||||
@@ -6235,10 +6229,13 @@ def test_start_branch_then(
|
||||
tool_two_graph.add_node("tool_two_slow", tool_two_slow)
|
||||
tool_two_graph.add_node("tool_two_fast", tool_two_fast)
|
||||
tool_two_graph.set_conditional_entry_point(
|
||||
lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast", then=END
|
||||
lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast",
|
||||
then=END,
|
||||
path_map=["tool_two_slow", "tool_two_fast"],
|
||||
)
|
||||
tool_two = tool_two_graph.compile()
|
||||
assert tool_two.get_graph().draw_mermaid() == snapshot
|
||||
if checkpointer_name == "memory":
|
||||
assert tool_two.get_graph().draw_mermaid() == snapshot
|
||||
|
||||
assert tool_two.invoke({"my_key": "value", "market": "DE"}) == {
|
||||
"my_key": "value slow",
|
||||
@@ -6517,6 +6514,7 @@ def test_branch_then(
|
||||
tool_two_graph.add_conditional_edges(
|
||||
source="prepare",
|
||||
path=lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast",
|
||||
path_map=["tool_two_slow", "tool_two_fast"],
|
||||
then="finish",
|
||||
)
|
||||
tool_two_graph.add_node("prepare", lambda s: {"my_key": " prepared"})
|
||||
@@ -6524,8 +6522,10 @@ def test_branch_then(
|
||||
tool_two_graph.add_node("tool_two_fast", lambda s: {"my_key": " fast"})
|
||||
tool_two_graph.add_node("finish", lambda s: {"my_key": " finished"})
|
||||
tool_two = tool_two_graph.compile()
|
||||
assert tool_two.get_graph().draw_mermaid(with_styles=False) == snapshot
|
||||
assert tool_two.get_graph().draw_mermaid() == snapshot
|
||||
|
||||
if checkpointer_name == "memory":
|
||||
assert tool_two.get_graph().draw_mermaid(with_styles=False) == snapshot
|
||||
assert tool_two.get_graph().draw_mermaid() == snapshot
|
||||
|
||||
assert tool_two.invoke({"my_key": "value", "market": "DE"}, debug=1) == {
|
||||
"my_key": "value prepared slow finished",
|
||||
@@ -9857,7 +9857,9 @@ def test_send_react_interrupt_control(
|
||||
builder.add_node(foo)
|
||||
builder.add_edge(START, "agent")
|
||||
graph = builder.compile()
|
||||
assert graph.get_graph().draw_mermaid() == snapshot
|
||||
|
||||
if checkpointer_name == "memory":
|
||||
assert graph.get_graph().draw_mermaid() == snapshot
|
||||
|
||||
assert graph.invoke({"messages": [HumanMessage("hello")]}) == {
|
||||
"messages": [
|
||||
@@ -10188,12 +10190,17 @@ def test_weather_subgraph(
|
||||
graph.add_node(normal_llm_node)
|
||||
graph.add_node("weather_graph", weather_graph)
|
||||
graph.add_edge(START, "router_node")
|
||||
graph.add_conditional_edges("router_node", route_after_prediction)
|
||||
graph.add_conditional_edges(
|
||||
"router_node",
|
||||
route_after_prediction,
|
||||
path_map=["weather_graph", "normal_llm_node"],
|
||||
)
|
||||
graph.add_edge("normal_llm_node", END)
|
||||
graph.add_edge("weather_graph", END)
|
||||
graph = graph.compile(checkpointer=checkpointer)
|
||||
|
||||
assert graph.get_graph(xray=1).draw_mermaid() == snapshot
|
||||
if checkpointer_name == "memory":
|
||||
assert graph.get_graph(xray=1).draw_mermaid() == snapshot
|
||||
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
thread2 = {"configurable": {"thread_id": "2"}}
|
||||
|
||||
@@ -7041,7 +7041,11 @@ async def test_weather_subgraph(
|
||||
graph.add_node(normal_llm_node)
|
||||
graph.add_node("weather_graph", weather_graph)
|
||||
graph.add_edge(START, "router_node")
|
||||
graph.add_conditional_edges("router_node", route_after_prediction)
|
||||
graph.add_conditional_edges(
|
||||
"router_node",
|
||||
route_after_prediction,
|
||||
path_map=["weather_graph", "normal_llm_node"],
|
||||
)
|
||||
graph.add_edge("normal_llm_node", END)
|
||||
graph.add_edge("weather_graph", END)
|
||||
|
||||
@@ -7051,8 +7055,6 @@ async def test_weather_subgraph(
|
||||
async with awith_checkpointer(checkpointer_name) as checkpointer:
|
||||
graph = graph.compile(checkpointer=checkpointer)
|
||||
|
||||
assert graph.get_graph(xray=1).draw_mermaid() == snapshot
|
||||
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
thread2 = {"configurable": {"thread_id": "2"}}
|
||||
inputs = {"messages": [{"role": "user", "content": "what's the weather in sf"}]}
|
||||
|
||||
@@ -1,14 +1,9 @@
|
||||
import datetime
|
||||
import decimal
|
||||
import enum
|
||||
import functools
|
||||
import gc
|
||||
import ipaddress
|
||||
import json
|
||||
import logging
|
||||
import operator
|
||||
import pathlib
|
||||
import re
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
@@ -18,7 +13,6 @@ from collections.abc import Generator, Iterator, Sequence
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from random import randrange
|
||||
from typing import (
|
||||
Annotated,
|
||||
@@ -2148,7 +2142,7 @@ def test_conditional_entrypoint_to_multiple_state_graph(
|
||||
|
||||
workflow.add_node("get_weather", get_weather)
|
||||
workflow.add_edge("get_weather", END)
|
||||
workflow.set_conditional_entry_point(continue_to_weather)
|
||||
workflow.set_conditional_entry_point(continue_to_weather, path_map=["get_weather"])
|
||||
|
||||
app = workflow.compile()
|
||||
|
||||
@@ -2415,7 +2409,8 @@ def test_in_one_fan_out_state_graph_waiting_edge(
|
||||
|
||||
app = workflow.compile()
|
||||
|
||||
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
|
||||
if checkpointer_name == "memory":
|
||||
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
|
||||
|
||||
assert app.invoke({"query": "what is weather in sf"}) == {
|
||||
"query": "analyzed: query: what is weather in sf",
|
||||
@@ -2561,7 +2556,8 @@ def test_in_one_fan_out_state_graph_waiting_edge_via_branch(
|
||||
|
||||
app = workflow.compile()
|
||||
|
||||
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
|
||||
if checkpointer_name == "memory":
|
||||
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
|
||||
|
||||
assert app.invoke({"query": "what is weather in sf"}, debug=True) == {
|
||||
"query": "analyzed: query: what is weather in sf",
|
||||
@@ -2711,9 +2707,10 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1(
|
||||
|
||||
app = workflow.compile()
|
||||
|
||||
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
|
||||
assert app.get_input_jsonschema() == snapshot
|
||||
assert app.get_output_jsonschema() == snapshot
|
||||
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": {}})
|
||||
@@ -2901,7 +2898,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2(
|
||||
|
||||
app = workflow.compile()
|
||||
|
||||
if SHOULD_CHECK_SNAPSHOTS:
|
||||
if SHOULD_CHECK_SNAPSHOTS and checkpointer_name == "memory":
|
||||
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
|
||||
assert app.get_input_schema().model_json_schema() == snapshot
|
||||
assert app.get_output_schema().model_json_schema() == snapshot
|
||||
@@ -2965,8 +2962,6 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2(
|
||||
|
||||
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
|
||||
def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_input(
|
||||
snapshot: SnapshotAssertion,
|
||||
mocker: MockerFixture,
|
||||
request: pytest.FixtureRequest,
|
||||
checkpointer_name: str,
|
||||
) -> None:
|
||||
@@ -3095,329 +3090,6 @@ 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
|
||||
# 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
|
||||
|
||||
# For constrained types
|
||||
PositiveInt = Annotated[int, Field(gt=0)]
|
||||
NonNegativeFloat = Annotated[float, Field(ge=0)]
|
||||
|
||||
# Enum type
|
||||
class UserRole(Enum):
|
||||
ADMIN = "admin"
|
||||
USER = "user"
|
||||
GUEST = "guest"
|
||||
|
||||
# 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
|
||||
|
||||
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)
|
||||
|
||||
class State(BaseModel):
|
||||
# Basic nested model tests
|
||||
top_level: str
|
||||
auuid: uuid.UUID
|
||||
nested: NestedModel
|
||||
optional_nested: Annotated[Optional[NestedModel], lambda x, y: y, "Foo"]
|
||||
dict_nested: dict[str, NestedModel]
|
||||
simple_str_list: list[str]
|
||||
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
|
||||
|
||||
# Rich type adapters
|
||||
ip_address: ipaddress.IPv4Address
|
||||
ip_address_v6: ipaddress.IPv6Address
|
||||
amount: decimal.Decimal
|
||||
file_path: pathlib.Path
|
||||
timestamp: datetime.datetime
|
||||
date_only: datetime.date
|
||||
time_only: datetime.time
|
||||
duration: datetime.timedelta
|
||||
immutable_set: frozenset[int]
|
||||
binary_data: bytes
|
||||
pattern: re.Pattern
|
||||
secret: SecretStr
|
||||
file_size: ByteSize
|
||||
|
||||
# Constrained types
|
||||
positive_value: PositiveInt
|
||||
non_negative: NonNegativeFloat
|
||||
limited_string: constr(min_length=3, max_length=10)
|
||||
bounded_int: conint(ge=10, le=100)
|
||||
restricted_float: confloat(gt=0, lt=1)
|
||||
required_list: conlist_type
|
||||
|
||||
# Enum & Literal
|
||||
role: UserRole
|
||||
status: Literal["active", "inactive", "pending"]
|
||||
|
||||
# Annotated & NewType
|
||||
validated_age: Annotated[int, Field(gt=0, lt=120)]
|
||||
|
||||
# Generic containers with validators
|
||||
decimal_list: list[decimal.Decimal]
|
||||
id_tuple: tuple[uuid.UUID, uuid.UUID]
|
||||
|
||||
inputs = {
|
||||
# Basic nested models
|
||||
"top_level": "initial",
|
||||
"auuid": str(uuid.uuid4()),
|
||||
"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"}]],
|
||||
"simple_str_list": ["siss", "boom", "bah"],
|
||||
"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
|
||||
},
|
||||
},
|
||||
# Rich type adapters
|
||||
"ip_address": "192.168.1.1",
|
||||
"ip_address_v6": "2001:db8::1",
|
||||
"amount": "123.45",
|
||||
"file_path": "/tmp/test.txt",
|
||||
"timestamp": "2025-04-07T10:58:04",
|
||||
"date_only": "2025-04-07",
|
||||
"time_only": "10:58:04",
|
||||
"duration": 3600, # seconds
|
||||
"immutable_set": [1, 2, 3, 4],
|
||||
"binary_data": b"hello world",
|
||||
"pattern": "^test$",
|
||||
"secret": "password123",
|
||||
"file_size": 1024,
|
||||
# Constrained types
|
||||
"positive_value": 42,
|
||||
"non_negative": 0.0,
|
||||
"limited_string": "test",
|
||||
"bounded_int": 50,
|
||||
"restricted_float": 0.5,
|
||||
"required_list": [10, 20, 30],
|
||||
# Enum & Literal
|
||||
"role": "admin",
|
||||
"status": "active",
|
||||
# Annotated & NewType
|
||||
"validated_age": 30,
|
||||
# Generic containers with validators
|
||||
"decimal_list": ["10.5", "20.75", "30.25"],
|
||||
"id_tuple": [str(uuid.uuid4()), str(uuid.uuid4())],
|
||||
}
|
||||
|
||||
update = {"top_level": "updated", "nested": {"value": 100, "name": "updated"}}
|
||||
|
||||
expected = State(**inputs)
|
||||
|
||||
def node_fn(state: State) -> dict:
|
||||
# Basic assertions
|
||||
assert isinstance(state.auuid, uuid.UUID)
|
||||
assert state == expected
|
||||
|
||||
# Rich type assertions
|
||||
assert isinstance(state.ip_address, ipaddress.IPv4Address)
|
||||
assert isinstance(state.ip_address_v6, ipaddress.IPv6Address)
|
||||
assert isinstance(state.amount, decimal.Decimal)
|
||||
assert isinstance(state.file_path, pathlib.Path)
|
||||
assert isinstance(state.timestamp, datetime.datetime)
|
||||
assert isinstance(state.date_only, datetime.date)
|
||||
assert isinstance(state.time_only, datetime.time)
|
||||
assert isinstance(state.duration, datetime.timedelta)
|
||||
assert isinstance(state.immutable_set, frozenset)
|
||||
assert isinstance(state.binary_data, bytes)
|
||||
assert isinstance(state.pattern, re.Pattern)
|
||||
|
||||
# Constrained types
|
||||
assert state.positive_value > 0
|
||||
assert state.non_negative >= 0
|
||||
assert 3 <= len(state.limited_string) <= 10
|
||||
assert 10 <= state.bounded_int <= 100
|
||||
assert 0 < state.restricted_float < 1
|
||||
assert 2 <= len(state.required_list) <= 5
|
||||
|
||||
# Enum & Literal
|
||||
assert state.role == UserRole.ADMIN
|
||||
assert state.status == "active"
|
||||
|
||||
# Annotated
|
||||
assert 0 < state.validated_age < 120
|
||||
|
||||
# Generic containers
|
||||
assert len(state.decimal_list) == 3
|
||||
assert len(state.id_tuple) == 2
|
||||
|
||||
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())
|
||||
|
||||
|
||||
def test_pydantic_state_field_validator():
|
||||
from pydantic import BaseModel, field_validator, model_validator
|
||||
|
||||
class State(BaseModel):
|
||||
name: str
|
||||
text: str = ""
|
||||
only_root: int = 13
|
||||
|
||||
@field_validator("name", mode="after")
|
||||
@classmethod
|
||||
def validate_name(cls, value):
|
||||
if value[0].islower():
|
||||
raise ValueError("Name must start with a capital letter")
|
||||
return "Validated " + value
|
||||
|
||||
@model_validator(mode="before")
|
||||
@classmethod
|
||||
def validate_amodel(cls, values: "State"):
|
||||
return values | {"only_root": 392}
|
||||
|
||||
input_state = {"name": "John"}
|
||||
|
||||
def process_node(state: State):
|
||||
assert State.model_validate(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!"
|
||||
|
||||
|
||||
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!"
|
||||
|
||||
|
||||
@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
|
||||
@@ -4799,7 +4471,9 @@ def test_xray_lance(snapshot: SnapshotAssertion):
|
||||
# Flow
|
||||
interview_builder.add_edge(START, "ask_question")
|
||||
interview_builder.add_edge("ask_question", "answer_question")
|
||||
interview_builder.add_conditional_edges("answer_question", route_messages)
|
||||
interview_builder.add_conditional_edges(
|
||||
"answer_question", route_messages, ["ask_question", END]
|
||||
)
|
||||
|
||||
# Set up memory
|
||||
memory = InMemorySaver()
|
||||
@@ -7331,6 +7005,8 @@ def test_node_destinations() -> None:
|
||||
Edge(source="__start__", target="child", data=None, conditional=False),
|
||||
Edge(source="child", target="node_b", data=None, conditional=True),
|
||||
Edge(source="child", target="node_c", data=None, conditional=True),
|
||||
Edge(source="node_b", target="__end__", data=None, conditional=False),
|
||||
Edge(source="node_c", target="__end__", data=None, conditional=False),
|
||||
] == graph.edges
|
||||
|
||||
# destinations w/ dicts
|
||||
@@ -7349,6 +7025,8 @@ def test_node_destinations() -> None:
|
||||
Edge(source="__start__", target="child", data=None, conditional=False),
|
||||
Edge(source="child", target="node_b", data="foo", conditional=True),
|
||||
Edge(source="child", target="node_c", data="bar", conditional=True),
|
||||
Edge(source="node_b", target="__end__", data=None, conditional=False),
|
||||
Edge(source="node_c", target="__end__", data=None, conditional=False),
|
||||
] == graph.edges
|
||||
|
||||
|
||||
|
||||
@@ -3605,7 +3605,8 @@ async def test_send_react_interrupt_control(
|
||||
builder.add_node(foo)
|
||||
builder.add_edge(START, "agent")
|
||||
graph = builder.compile()
|
||||
assert graph.get_graph().draw_mermaid() == snapshot
|
||||
if checkpointer_name == "memory":
|
||||
assert graph.get_graph().draw_mermaid() == snapshot
|
||||
|
||||
assert await graph.ainvoke({"messages": [HumanMessage("hello")]}) == {
|
||||
"messages": [
|
||||
@@ -3923,22 +3924,29 @@ async def test_max_concurrency_control(checkpointer_name: str) -> None:
|
||||
builder.add_edge(START, "1")
|
||||
graph = builder.compile()
|
||||
|
||||
assert (
|
||||
graph.get_graph().draw_mermaid()
|
||||
== """%%{init: {'flowchart': {'curve': 'linear'}}}%%
|
||||
if checkpointer_name == "memory":
|
||||
assert (
|
||||
graph.get_graph().draw_mermaid()
|
||||
== """---
|
||||
config:
|
||||
flowchart:
|
||||
curve: linear
|
||||
---
|
||||
graph TD;
|
||||
__start__([<p>__start__</p>]):::first
|
||||
1(1)
|
||||
2(2)
|
||||
3([3]):::last
|
||||
__start__ --> 1;
|
||||
3(3)
|
||||
__end__([<p>__end__</p>]):::last
|
||||
1 -.-> 2;
|
||||
2 -.-> 3;
|
||||
__start__ --> 1;
|
||||
3 --> __end__;
|
||||
classDef default fill:#f2f0ff,line-height:1.2
|
||||
classDef first fill-opacity:0
|
||||
classDef last fill:#bfb6fc
|
||||
"""
|
||||
)
|
||||
)
|
||||
|
||||
assert await graph.ainvoke(["0"], debug=True) == ["0", "1", *range(100), "3"]
|
||||
assert node2_max_currently == 100
|
||||
@@ -4975,7 +4983,7 @@ async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydant
|
||||
|
||||
app = workflow.compile()
|
||||
|
||||
if SHOULD_CHECK_SNAPSHOTS:
|
||||
if SHOULD_CHECK_SNAPSHOTS and checkpointer_name == "memory":
|
||||
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
|
||||
assert app.get_input_schema().model_json_schema() == snapshot
|
||||
assert app.get_output_schema().model_json_schema() == snapshot
|
||||
|
||||
@@ -1,14 +1,26 @@
|
||||
import datetime
|
||||
import decimal
|
||||
import ipaddress
|
||||
import pathlib
|
||||
import re
|
||||
import sys
|
||||
import typing
|
||||
import uuid
|
||||
from enum import Enum
|
||||
from typing import Annotated, List, Literal, Optional, Union
|
||||
|
||||
import pydantic
|
||||
import typing_extensions
|
||||
import pytest
|
||||
|
||||
from langgraph.constants import END, START
|
||||
from langgraph.graph.state import StateGraph
|
||||
from langgraph.utils.pydantic import is_supported_by_pydantic
|
||||
|
||||
|
||||
def test_is_supported_by_pydantic() -> None:
|
||||
"""Test if types are supported by pydantic."""
|
||||
import typing
|
||||
|
||||
import pydantic
|
||||
import typing_extensions
|
||||
|
||||
class TypedDictExtensions(typing_extensions.TypedDict):
|
||||
x: int
|
||||
@@ -41,3 +53,325 @@ def test_is_supported_by_pydantic() -> None:
|
||||
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:
|
||||
"""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
|
||||
|
||||
# For constrained types
|
||||
PositiveInt = Annotated[int, Field(gt=0)]
|
||||
NonNegativeFloat = Annotated[float, Field(ge=0)]
|
||||
|
||||
# Enum type
|
||||
class UserRole(Enum):
|
||||
ADMIN = "admin"
|
||||
USER = "user"
|
||||
GUEST = "guest"
|
||||
|
||||
# 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
|
||||
|
||||
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)
|
||||
|
||||
class State(BaseModel):
|
||||
# Basic nested model tests
|
||||
top_level: str
|
||||
auuid: uuid.UUID
|
||||
nested: NestedModel
|
||||
optional_nested: Annotated[Optional[NestedModel], lambda x, y: y, "Foo"]
|
||||
dict_nested: dict[str, NestedModel]
|
||||
simple_str_list: list[str]
|
||||
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
|
||||
|
||||
# Rich type adapters
|
||||
ip_address: ipaddress.IPv4Address
|
||||
ip_address_v6: ipaddress.IPv6Address
|
||||
amount: decimal.Decimal
|
||||
file_path: pathlib.Path
|
||||
timestamp: datetime.datetime
|
||||
date_only: datetime.date
|
||||
time_only: datetime.time
|
||||
duration: datetime.timedelta
|
||||
immutable_set: frozenset[int]
|
||||
binary_data: bytes
|
||||
pattern: re.Pattern
|
||||
secret: SecretStr
|
||||
file_size: ByteSize
|
||||
|
||||
# Constrained types
|
||||
positive_value: PositiveInt
|
||||
non_negative: NonNegativeFloat
|
||||
limited_string: constr(min_length=3, max_length=10)
|
||||
bounded_int: conint(ge=10, le=100)
|
||||
restricted_float: confloat(gt=0, lt=1)
|
||||
required_list: conlist_type
|
||||
|
||||
# Enum & Literal
|
||||
role: UserRole
|
||||
status: Literal["active", "inactive", "pending"]
|
||||
|
||||
# Annotated & NewType
|
||||
validated_age: Annotated[int, Field(gt=0, lt=120)]
|
||||
|
||||
# Generic containers with validators
|
||||
decimal_list: List[decimal.Decimal]
|
||||
id_tuple: tuple[uuid.UUID, uuid.UUID]
|
||||
|
||||
inputs = {
|
||||
# Basic nested models
|
||||
"top_level": "initial",
|
||||
"auuid": str(uuid.uuid4()),
|
||||
"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"}]],
|
||||
"simple_str_list": ["siss", "boom", "bah"],
|
||||
"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
|
||||
},
|
||||
},
|
||||
# Rich type adapters
|
||||
"ip_address": "192.168.1.1",
|
||||
"ip_address_v6": "2001:db8::1",
|
||||
"amount": "123.45",
|
||||
"file_path": "/tmp/test.txt",
|
||||
"timestamp": "2025-04-07T10:58:04",
|
||||
"date_only": "2025-04-07",
|
||||
"time_only": "10:58:04",
|
||||
"duration": 3600, # seconds
|
||||
"immutable_set": [1, 2, 3, 4],
|
||||
"binary_data": b"hello world",
|
||||
"pattern": "^test$",
|
||||
"secret": "password123",
|
||||
"file_size": 1024,
|
||||
# Constrained types
|
||||
"positive_value": 42,
|
||||
"non_negative": 0.0,
|
||||
"limited_string": "test",
|
||||
"bounded_int": 50,
|
||||
"restricted_float": 0.5,
|
||||
"required_list": [10, 20, 30],
|
||||
# Enum & Literal
|
||||
"role": "admin",
|
||||
"status": "active",
|
||||
# Annotated & NewType
|
||||
"validated_age": 30,
|
||||
# Generic containers with validators
|
||||
"decimal_list": ["10.5", "20.75", "30.25"],
|
||||
"id_tuple": [str(uuid.uuid4()), str(uuid.uuid4())],
|
||||
}
|
||||
|
||||
update = {"top_level": "updated", "nested": {"value": 100, "name": "updated"}}
|
||||
|
||||
expected = State(**inputs)
|
||||
|
||||
def node_fn(state: State) -> dict:
|
||||
# Basic assertions
|
||||
assert isinstance(state.auuid, uuid.UUID)
|
||||
assert state == expected
|
||||
|
||||
# Rich type assertions
|
||||
assert isinstance(state.ip_address, ipaddress.IPv4Address)
|
||||
assert isinstance(state.ip_address_v6, ipaddress.IPv6Address)
|
||||
assert isinstance(state.amount, decimal.Decimal)
|
||||
assert isinstance(state.file_path, pathlib.Path)
|
||||
assert isinstance(state.timestamp, datetime.datetime)
|
||||
assert isinstance(state.date_only, datetime.date)
|
||||
assert isinstance(state.time_only, datetime.time)
|
||||
assert isinstance(state.duration, datetime.timedelta)
|
||||
assert isinstance(state.immutable_set, frozenset)
|
||||
assert isinstance(state.binary_data, bytes)
|
||||
assert isinstance(state.pattern, re.Pattern)
|
||||
|
||||
# Constrained types
|
||||
assert state.positive_value > 0
|
||||
assert state.non_negative >= 0
|
||||
assert 3 <= len(state.limited_string) <= 10
|
||||
assert 10 <= state.bounded_int <= 100
|
||||
assert 0 < state.restricted_float < 1
|
||||
assert 2 <= len(state.required_list) <= 5
|
||||
|
||||
# Enum & Literal
|
||||
assert state.role == UserRole.ADMIN
|
||||
assert state.status == "active"
|
||||
|
||||
# Annotated
|
||||
assert 0 < state.validated_age < 120
|
||||
|
||||
# Generic containers
|
||||
assert len(state.decimal_list) == 3
|
||||
assert len(state.id_tuple) == 2
|
||||
|
||||
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())
|
||||
|
||||
|
||||
def test_pydantic_state_field_validator():
|
||||
from pydantic import BaseModel, field_validator, model_validator
|
||||
|
||||
class State(BaseModel):
|
||||
name: str
|
||||
text: str = ""
|
||||
only_root: int = 13
|
||||
|
||||
@field_validator("name", mode="after")
|
||||
@classmethod
|
||||
def validate_name(cls, value):
|
||||
if value[0].islower():
|
||||
raise ValueError("Name must start with a capital letter")
|
||||
return "Validated " + value
|
||||
|
||||
@model_validator(mode="before")
|
||||
@classmethod
|
||||
def validate_amodel(cls, values: "State"):
|
||||
return values | {"only_root": 392}
|
||||
|
||||
input_state = {"name": "John"}
|
||||
|
||||
def process_node(state: State):
|
||||
assert State.model_validate(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!"
|
||||
|
||||
|
||||
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!"
|
||||
|
||||
Reference in New Issue
Block a user