diff --git a/libs/checkpoint-postgres/README.md b/libs/checkpoint-postgres/README.md index 24652a2b2..cf6beabac 100644 --- a/libs/checkpoint-postgres/README.md +++ b/libs/checkpoint-postgres/README.md @@ -2,6 +2,10 @@ Implementation of LangGraph CheckpointSaver that uses Postgres. +## Dependencies + +By default `langgraph-checkpoint-postgres` installs `psycopg` (Psycopg 3) without any extras. However, you can choose a specific installation that best suits your needs [here](https://www.psycopg.org/psycopg3/docs/basic/install.html) (for example, `psycopg[binary]`). + ## Usage > [!IMPORTANT] diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py index cb4a79c35..3ce8e880c 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py @@ -1,12 +1,13 @@ import threading from contextlib import contextmanager -from typing import Any, Iterator, List, Optional +from typing import Any, Iterator, List, Optional, Union from langchain_core.runnables import RunnableConfig from psycopg import Connection, Cursor, Pipeline from psycopg.errors import UndefinedTable from psycopg.rows import dict_row from psycopg.types.json import Jsonb +from psycopg_pool import ConnectionPool from langgraph.checkpoint.base import ( ChannelVersions, @@ -21,16 +22,32 @@ from langgraph.checkpoint.postgres.base import ( from langgraph.checkpoint.serde.base import SerializerProtocol +@contextmanager +def _get_connection(conn: Union[Connection, ConnectionPool]) -> Iterator[Connection]: + if isinstance(conn, Connection): + yield conn + elif isinstance(conn, ConnectionPool): + with conn.connection() as conn: + yield conn + else: + raise TypeError(f"Invalid connection type: {type(conn)}") + + class PostgresSaver(BasePostgresSaver): lock: threading.Lock def __init__( self, - conn: Connection, + conn: Union[Connection, ConnectionPool], pipe: Optional[Pipeline] = None, serde: Optional[SerializerProtocol] = None, ) -> None: super().__init__(serde=serde) + if isinstance(conn, ConnectionPool) and pipe is not None: + raise ValueError( + "Pipeline should be used only with a single Connection, not ConnectionPool." + ) + self.conn = conn self.pipe = pipe self.lock = threading.Lock() @@ -65,22 +82,21 @@ class PostgresSaver(BasePostgresSaver): already exist and runs database migrations. It MUST be called directly by the user the first time checkpointer is used. """ - with self.lock: - with self.conn.cursor(binary=True, row_factory=dict_row) as cur: - try: - version = cur.execute( - "SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1" - ).fetchone()["v"] - except UndefinedTable: - version = -1 - for v, migration in zip( - range(version + 1, len(self.MIGRATIONS)), - self.MIGRATIONS[version + 1 :], - ): - cur.execute(migration) - cur.execute(f"INSERT INTO checkpoint_migrations (v) VALUES ({v})") - if self.pipe: - self.pipe.sync() + with self._cursor() as cur: + try: + version = cur.execute( + "SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1" + ).fetchone()["v"] + except UndefinedTable: + version = -1 + for v, migration in zip( + range(version + 1, len(self.MIGRATIONS)), + self.MIGRATIONS[version + 1 :], + ): + cur.execute(migration) + cur.execute(f"INSERT INTO checkpoint_migrations (v) VALUES ({v})") + if self.pipe: + self.pipe.sync() def list( self, @@ -333,23 +349,24 @@ class PostgresSaver(BasePostgresSaver): @contextmanager def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor]: - if self.pipe: - # a connection in pipeline mode can be used concurrently - # in multiple threads/coroutines, but only one cursor can be - # used at a time - try: - with self.conn.cursor(binary=True, row_factory=dict_row) as cur: + with _get_connection(self.conn) as conn: + if self.pipe: + # a connection in pipeline mode can be used concurrently + # in multiple threads/coroutines, but only one cursor can be + # used at a time + try: + with conn.cursor(binary=True, row_factory=dict_row) as cur: + yield cur + finally: + if pipeline: + self.pipe.sync() + elif pipeline: + # a connection not in pipeline mode can only be used by one + # thread/coroutine at a time, so we acquire a lock + with self.lock, conn.pipeline(), conn.cursor( + binary=True, row_factory=dict_row + ) as cur: + yield cur + else: + with self.lock, conn.cursor(binary=True, row_factory=dict_row) as cur: yield cur - finally: - if pipeline: - self.pipe.sync() - elif pipeline: - # a connection not in pipeline mode can only be used by one - # thread/coroutine at a time, so we acquire a lock - with self.lock, self.conn.pipeline(), self.conn.cursor( - binary=True, row_factory=dict_row - ) as cur: - yield cur - else: - with self.lock, self.conn.cursor(binary=True, row_factory=dict_row) as cur: - yield cur diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index 569159d91..0db61a443 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -1,12 +1,13 @@ import asyncio from contextlib import asynccontextmanager -from typing import Any, AsyncIterator, Optional +from typing import Any, AsyncIterator, Optional, Union from langchain_core.runnables import RunnableConfig from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline from psycopg.errors import UndefinedTable from psycopg.rows import dict_row from psycopg.types.json import Jsonb +from psycopg_pool import AsyncConnectionPool from langgraph.checkpoint.base import ( ChannelVersions, @@ -19,16 +20,34 @@ from langgraph.checkpoint.postgres.base import BasePostgresSaver from langgraph.checkpoint.serde.base import SerializerProtocol +@asynccontextmanager +async def _get_connection( + conn: Union[AsyncConnection, AsyncConnectionPool], +) -> AsyncIterator[AsyncConnection]: + if isinstance(conn, AsyncConnection): + yield conn + elif isinstance(conn, AsyncConnectionPool): + async with conn.connection() as conn: + yield conn + else: + raise TypeError(f"Invalid connection type: {type(conn)}") + + class AsyncPostgresSaver(BasePostgresSaver): lock: asyncio.Lock def __init__( self, - conn: AsyncConnection, + conn: Union[AsyncConnection, AsyncConnectionPool], pipe: Optional[AsyncPipeline] = None, serde: Optional[SerializerProtocol] = None, ) -> None: super().__init__(serde=serde) + if isinstance(conn, AsyncConnectionPool) and pipe is not None: + raise ValueError( + "Pipeline should be used only with a single AsyncConnection, not AsyncConnectionPool." + ) + self.conn = conn self.pipe = pipe self.lock = asyncio.Lock() @@ -63,25 +82,22 @@ class AsyncPostgresSaver(BasePostgresSaver): already exist and runs database migrations. It MUST be called directly by the user the first time checkpointer is used. """ - async with self.lock: - async with self.conn.cursor(binary=True, row_factory=dict_row) as cur: - try: - results = await cur.execute( - "SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1" - ) - version = (await results.fetchone())["v"] - except UndefinedTable: - version = -1 - for v, migration in zip( - range(version + 1, len(self.MIGRATIONS)), - self.MIGRATIONS[version + 1 :], - ): - await cur.execute(migration) - await cur.execute( - f"INSERT INTO checkpoint_migrations (v) VALUES ({v})" - ) - if self.pipe: - await self.pipe.sync() + async with self._cursor() as cur: + try: + results = await cur.execute( + "SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1" + ) + version = (await results.fetchone())["v"] + except UndefinedTable: + version = -1 + for v, migration in zip( + range(version + 1, len(self.MIGRATIONS)), + self.MIGRATIONS[version + 1 :], + ): + await cur.execute(migration) + await cur.execute(f"INSERT INTO checkpoint_migrations (v) VALUES ({v})") + if self.pipe: + await self.pipe.sync() async def alist( self, @@ -290,25 +306,26 @@ class AsyncPostgresSaver(BasePostgresSaver): @asynccontextmanager async def _cursor(self, *, pipeline: bool = False) -> AsyncIterator[AsyncCursor]: - if self.pipe: - # a connection in pipeline mode can be used concurrently - # in multiple threads/coroutines, but only one cursor can be - # used at a time - try: - async with self.conn.cursor(binary=True, row_factory=dict_row) as cur: + async with _get_connection(self.conn) as conn: + if self.pipe: + # a connection in pipeline mode can be used concurrently + # in multiple threads/coroutines, but only one cursor can be + # used at a time + try: + async with conn.cursor(binary=True, row_factory=dict_row) as cur: + yield cur + finally: + if pipeline: + await self.pipe.sync() + elif pipeline: + # a connection not in pipeline mode can only be used by one + # thread/coroutine at a time, so we acquire a lock + async with self.lock, conn.pipeline(), conn.cursor( + binary=True, row_factory=dict_row + ) as cur: + yield cur + else: + async with self.lock, conn.cursor( + binary=True, row_factory=dict_row + ) as cur: yield cur - finally: - if pipeline: - await self.pipe.sync() - elif pipeline: - # a connection not in pipeline mode can only be used by one - # thread/coroutine at a time, so we acquire a lock - async with self.lock, self.conn.pipeline(), self.conn.cursor( - binary=True, row_factory=dict_row - ) as cur: - yield cur - else: - async with self.lock, self.conn.cursor( - binary=True, row_factory=dict_row - ) as cur: - yield cur diff --git a/libs/checkpoint-postgres/poetry.lock b/libs/checkpoint-postgres/poetry.lock index 813bdf9b0..0c171e445 100644 --- a/libs/checkpoint-postgres/poetry.lock +++ b/libs/checkpoint-postgres/poetry.lock @@ -266,7 +266,7 @@ tenacity = ">=8.1.0,<8.4.0 || >8.4.0,<9.0.0" [[package]] name = "langgraph-checkpoint" -version = "1.0.1" +version = "1.0.6" description = "Library with base interfaces for LangGraph checkpoint savers." optional = false python-versions = "^3.9.0,<4.0" @@ -969,4 +969,4 @@ watchmedo = ["PyYAML (>=3.10)"] [metadata] lock-version = "2.0" python-versions = "^3.9.0,<4.0" -content-hash = "422b6d716b86db072ea3a612287ad20ff5700c18f22d9e9d59cc4e198514519d" +content-hash = "cfa417cbaf126fa847242ef81e9af3e4d13e93f25631bb4c007cdc69926ff3b8" diff --git a/libs/checkpoint-postgres/pyproject.toml b/libs/checkpoint-postgres/pyproject.toml index d7f6899b0..ddf3a8107 100644 --- a/libs/checkpoint-postgres/pyproject.toml +++ b/libs/checkpoint-postgres/pyproject.toml @@ -12,7 +12,8 @@ packages = [{ include = "langgraph" }] python = "^3.9.0,<4.0" langgraph-checkpoint = "^1.0.1" orjson = ">=3.10.1" -psycopg = {extras = ["binary"], version = ">=3.1.19"} +psycopg = "^3.0.0" +psycopg-pool = "^3.0.0" [tool.poetry.group.dev.dependencies] ruff = "^0.1.4" @@ -23,7 +24,7 @@ pytest-asyncio = "^0.21.1" pytest-mock = "^3.11.1" pytest-watch = "^4.2.0" mypy = "^1.10.0" -psycopg-pool = "^3.2.2" +psycopg = {extras = ["binary"], version = ">=3.0.0"} langgraph-checkpoint = {path = "../checkpoint", develop = true} [tool.pytest.ini_options] diff --git a/libs/langgraph/poetry.lock b/libs/langgraph/poetry.lock index aa0054283..caa98200a 100644 --- a/libs/langgraph/poetry.lock +++ b/libs/langgraph/poetry.lock @@ -1832,7 +1832,7 @@ types-requests = ">=2.31.0.2,<3.0.0.0" [[package]] name = "langgraph-checkpoint" -version = "1.0.2" +version = "1.0.6" description = "Library with base interfaces for LangGraph checkpoint savers." optional = false python-versions = "^3.9.0,<4.0" @@ -1848,7 +1848,7 @@ url = "../checkpoint" [[package]] name = "langgraph-checkpoint-postgres" -version = "1.0.0" +version = "1.0.3" description = "Library with a Postgres implementation of LangGraph checkpoint saver." optional = false python-versions = "^3.9.0,<4.0" @@ -1858,7 +1858,8 @@ develop = true [package.dependencies] langgraph-checkpoint = "^1.0.1" orjson = ">=3.10.1" -psycopg = {version = ">=3.1.19", extras = ["binary"]} +psycopg = "^3.0.0" +psycopg-pool = "^3.0.0" [package.source] type = "directory" @@ -2638,6 +2639,20 @@ files = [ {file = "psycopg_binary-3.2.1-cp39-cp39-win_amd64.whl", hash = "sha256:921f0c7f39590763d64a619de84d1b142587acc70fd11cbb5ba8fa39786f3073"}, ] +[[package]] +name = "psycopg-pool" +version = "3.2.2" +description = "Connection Pool for Psycopg" +optional = false +python-versions = ">=3.8" +files = [ + {file = "psycopg_pool-3.2.2-py3-none-any.whl", hash = "sha256:273081d0fbfaced4f35e69200c89cb8fbddfe277c38cc86c235b90a2ec2c8153"}, + {file = "psycopg_pool-3.2.2.tar.gz", hash = "sha256:9e22c370045f6d7f2666a5ad1b0caf345f9f1912195b0b25d0d3bcc4f3a7389c"}, +] + +[package.dependencies] +typing-extensions = ">=4.4" + [[package]] name = "ptyprocess" version = "0.7.0" @@ -4294,4 +4309,4 @@ test = ["big-O", "importlib-resources", "jaraco.functools", "jaraco.itertools", [metadata] lock-version = "2.0" python-versions = ">=3.9.0,<4.0" -content-hash = "4ef9e25016072ce08554c8ff8d091104fb59da7cdab9a128312bd787e6f35146" +content-hash = "b4d234e851639ecb33b6507b5e396b8b87aae6395c327441d04a7195eae72582" diff --git a/libs/langgraph/pyproject.toml b/libs/langgraph/pyproject.toml index 35164c2f4..4e14ee41e 100644 --- a/libs/langgraph/pyproject.toml +++ b/libs/langgraph/pyproject.toml @@ -35,6 +35,7 @@ pytest-repeat = "^0.9.3" langgraph-checkpoint = {path = "../checkpoint", develop = true} langgraph-checkpoint-sqlite = {path = "../checkpoint-sqlite", develop = true} langgraph-checkpoint-postgres = {path = "../checkpoint-postgres", develop = true} +psycopg = {extras = ["binary"], version = ">=3.0.0"} [tool.poetry.group.dev] optional = true diff --git a/libs/langgraph/tests/__snapshots__/test_pregel.ambr b/libs/langgraph/tests/__snapshots__/test_pregel.ambr index 91096bca6..d273ea9a7 100644 --- a/libs/langgraph/tests/__snapshots__/test_pregel.ambr +++ b/libs/langgraph/tests/__snapshots__/test_pregel.ambr @@ -135,6 +135,40 @@ ''' # --- +# name: test_branch_then[postgres_pool] + ''' + graph TD; + __start__ --> prepare; + finish --> __end__; + prepare -.-> tool_two_slow; + tool_two_slow --> finish; + prepare -.-> tool_two_fast; + tool_two_fast --> finish; + + ''' +# --- +# name: test_branch_then[postgres_pool].1 + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__]):::first + prepare(prepare) + tool_two_slow(tool_two_slow) + tool_two_fast(tool_two_fast) + finish(finish) + __end__([__end__]):::last + __start__ --> prepare; + finish --> __end__; + prepare -.-> tool_two_slow; + tool_two_slow --> finish; + prepare -.-> tool_two_fast; + tool_two_fast --> finish; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- # name: test_branch_then[sqlite] ''' graph TD; @@ -1772,6 +1806,351 @@ ''' # --- +# name: test_conditional_graph[postgres_pool] + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": "__start__" + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": [ + "langchain", + "schema", + "runnable", + "RunnableAssign" + ], + "name": "agent" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "utils", + "RunnableCallable" + ], + "name": "tools" + }, + "metadata": { + "version": 2, + "variant": "b" + } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" + } + ], + "edges": [ + { + "source": "__start__", + "target": "agent" + }, + { + "source": "tools", + "target": "agent" + }, + { + "source": "agent", + "target": "tools", + "data": "continue", + "conditional": true + }, + { + "source": "agent", + "target": "__end__", + "data": "exit", + "conditional": true + } + ] + } + ''' +# --- +# name: test_conditional_graph[postgres_pool].1 + ''' + graph TD; + __start__ --> agent; + tools --> agent; + agent -.  continue  .-> tools; + agent -.  exit  .-> __end__; + + ''' +# --- +# name: test_conditional_graph[postgres_pool].2 + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__]):::first + agent(agent) + tools(tools
version = 2 + variant = b) + __end__([__end__]):::last + __start__ --> agent; + tools --> agent; + agent -.  continue  .-> tools; + agent -.  exit  .-> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_conditional_graph[postgres_pool].3 + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": "__start__" + }, + { + "id": 1, + "type": "schema", + "data": "ParallelInput" + }, + { + "id": 2, + "type": "schema", + "data": "ParallelOutput" + }, + { + "id": 3, + "type": "runnable", + "data": { + "id": [ + "langchain", + "prompts", + "prompt", + "PromptTemplate" + ], + "name": "PromptTemplate" + } + }, + { + "id": 4, + "type": "runnable", + "data": { + "id": [ + "langchain_core", + "language_models", + "fake", + "FakeStreamingListLLM" + ], + "name": "FakeStreamingListLLM" + } + }, + { + "id": 5, + "type": "runnable", + "data": { + "id": [ + "langchain_core", + "runnables", + "base", + "RunnableLambda" + ], + "name": "agent_parser" + } + }, + { + "id": 6, + "type": "runnable", + "data": { + "id": [ + "langchain", + "schema", + "runnable", + "RunnablePassthrough" + ], + "name": "Passthrough" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "utils", + "RunnableCallable" + ], + "name": "tools" + }, + "metadata": { + "version": 2, + "variant": "b" + } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" + } + ], + "edges": [ + { + "source": 3, + "target": 4 + }, + { + "source": 4, + "target": 5 + }, + { + "source": 1, + "target": 3 + }, + { + "source": 5, + "target": 2 + }, + { + "source": 1, + "target": 6 + }, + { + "source": 6, + "target": 2 + }, + { + "source": "__start__", + "target": 1 + }, + { + "source": "tools", + "target": 1 + }, + { + "source": 2, + "target": "tools", + "data": "continue", + "conditional": true + }, + { + "source": 2, + "target": "__end__", + "data": "exit", + "conditional": true + } + ] + } + ''' +# --- +# name: test_conditional_graph[postgres_pool].4 + ''' + graph TD; + PromptTemplate --> FakeStreamingListLLM; + FakeStreamingListLLM --> agent_parser; + Parallel_agent_outcome_Input --> PromptTemplate; + agent_parser --> Parallel_agent_outcome_Output; + Parallel_agent_outcome_Input --> Passthrough; + Passthrough --> Parallel_agent_outcome_Output; + __start__ --> Parallel_agent_outcome_Input; + tools --> Parallel_agent_outcome_Input; + Parallel_agent_outcome_Output -.  continue  .-> tools; + Parallel_agent_outcome_Output -.  exit  .-> __end__; + + ''' +# --- +# name: test_conditional_graph[postgres_pool].5 + dict({ + 'edges': list([ + dict({ + 'source': '__start__', + 'target': 'agent', + }), + dict({ + 'source': 'tools', + 'target': 'agent', + }), + dict({ + 'conditional': True, + 'data': 'continue', + 'source': 'agent', + 'target': 'tools', + }), + dict({ + 'conditional': True, + 'data': 'exit', + 'source': 'agent', + 'target': '__end__', + }), + ]), + 'nodes': list([ + dict({ + 'data': '__start__', + 'id': '__start__', + 'type': 'schema', + }), + dict({ + 'data': dict({ + 'id': list([ + 'langchain', + 'schema', + 'runnable', + 'RunnableAssign', + ]), + 'name': 'agent', + }), + 'id': 'agent', + 'metadata': dict({ + '__interrupt': 'after', + }), + 'type': 'runnable', + }), + dict({ + 'data': dict({ + 'id': list([ + 'langgraph', + 'utils', + 'RunnableCallable', + ]), + 'name': 'tools', + }), + 'id': 'tools', + 'metadata': dict({ + 'variant': 'b', + 'version': 2, + }), + 'type': 'runnable', + }), + dict({ + 'data': '__end__', + 'id': '__end__', + 'type': 'schema', + }), + ]), + }) +# --- +# name: test_conditional_graph[postgres_pool].6 + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__]):::first + agent(agent
__interrupt = after) + tools(tools
version = 2 + variant = b) + __end__([__end__]):::last + __start__ --> agent; + tools --> agent; + agent -.  continue  .-> tools; + agent -.  exit  .-> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- # name: test_conditional_graph[sqlite] ''' { @@ -2511,6 +2890,87 @@ ''' # --- +# name: test_conditional_state_graph[postgres_pool] + '{"title": "LangGraphInput", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}, "definitions": {"AgentAction": {"title": "AgentAction", "description": "Represents a request to execute an action by an agent.\\n\\nThe action consists of the name of the tool to execute and the input to pass\\nto the tool. The log is used to pass along extra information about the action.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "Final return value of an ActionAgent.\\n\\nAgents return an AgentFinish when they have reached a stopping condition.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}}}' +# --- +# name: test_conditional_state_graph[postgres_pool].1 + '{"title": "LangGraphOutput", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}, "definitions": {"AgentAction": {"title": "AgentAction", "description": "Represents a request to execute an action by an agent.\\n\\nThe action consists of the name of the tool to execute and the input to pass\\nto the tool. The log is used to pass along extra information about the action.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "Final return value of an ActionAgent.\\n\\nAgents return an AgentFinish when they have reached a stopping condition.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}}}' +# --- +# name: test_conditional_state_graph[postgres_pool].2 + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": "__start__" + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": [ + "langchain", + "schema", + "runnable", + "RunnableSequence" + ], + "name": "agent" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "utils", + "RunnableCallable" + ], + "name": "tools" + } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" + } + ], + "edges": [ + { + "source": "__start__", + "target": "agent" + }, + { + "source": "tools", + "target": "agent" + }, + { + "source": "agent", + "target": "tools", + "data": "continue", + "conditional": true + }, + { + "source": "agent", + "target": "__end__", + "data": "exit", + "conditional": true + } + ] + } + ''' +# --- +# name: test_conditional_state_graph[postgres_pool].3 + ''' + graph TD; + __start__ --> agent; + tools --> agent; + agent -.  continue  .-> tools; + agent -.  exit  .-> __end__; + + ''' +# --- # name: test_conditional_state_graph[sqlite] '{"title": "LangGraphInput", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}, "definitions": {"AgentAction": {"title": "AgentAction", "description": "Represents a request to execute an action by an agent.\\n\\nThe action consists of the name of the tool to execute and the input to pass\\nto the tool. The log is used to pass along extra information about the action.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "Final return value of an ActionAgent.\\n\\nAgents return an AgentFinish when they have reached a stopping condition.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}}}' # --- @@ -2662,6 +3122,19 @@ ''' # --- +# name: test_in_one_fan_out_state_graph_waiting_edge[postgres_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[sqlite] ''' graph TD; @@ -2968,6 +3441,76 @@ 'type': 'object', }) # --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[postgres_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_pydantic1[postgres_pool].1 + dict({ + 'definitions': dict({ + 'InnerObject': dict({ + 'properties': dict({ + 'yo': dict({ + 'title': 'Yo', + 'type': 'integer', + }), + }), + 'required': list([ + 'yo', + ]), + 'title': 'InnerObject', + 'type': 'object', + }), + }), + 'properties': dict({ + 'inner': dict({ + '$ref': '#/definitions/InnerObject', + }), + 'query': dict({ + 'title': 'Query', + 'type': 'string', + }), + }), + 'required': list([ + 'query', + 'inner', + ]), + 'title': 'Input', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[postgres_pool].2 + dict({ + 'properties': dict({ + 'answer': dict({ + 'title': 'Answer', + 'type': 'string', + }), + 'docs': dict({ + 'items': dict({ + 'type': 'string', + }), + 'title': 'Docs', + 'type': 'array', + }), + }), + 'required': list([ + 'answer', + 'docs', + ]), + 'title': 'Output', + 'type': 'object', + }) +# --- # name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[sqlite] ''' graph TD; @@ -3318,6 +3861,76 @@ 'type': 'object', }) # --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_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_pool].1 + dict({ + '$defs': dict({ + 'InnerObject': dict({ + 'properties': dict({ + 'yo': dict({ + 'title': 'Yo', + 'type': 'integer', + }), + }), + 'required': list([ + 'yo', + ]), + 'title': 'InnerObject', + 'type': 'object', + }), + }), + 'properties': dict({ + 'inner': dict({ + '$ref': '#/$defs/InnerObject', + }), + 'query': dict({ + 'title': 'Query', + 'type': 'string', + }), + }), + 'required': list([ + 'query', + 'inner', + ]), + 'title': 'Input', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_pool].2 + dict({ + 'properties': dict({ + 'answer': dict({ + 'title': 'Answer', + 'type': 'string', + }), + 'docs': dict({ + 'items': dict({ + 'type': 'string', + }), + 'title': 'Docs', + 'type': 'array', + }), + }), + 'required': list([ + 'answer', + 'docs', + ]), + 'title': 'Output', + 'type': 'object', + }) +# --- # name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[sqlite] ''' graph TD; @@ -3440,6 +4053,19 @@ ''' # --- +# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[postgres_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_via_branch[sqlite] ''' graph TD; @@ -3777,6 +4403,87 @@ ''' # --- +# name: test_message_graph[postgres_pool] + '{"title": "LangGraphInput", "type": "array", "items": {"anyOf": [{"$ref": "#/definitions/AIMessage"}, {"$ref": "#/definitions/HumanMessage"}, {"$ref": "#/definitions/ChatMessage"}, {"$ref": "#/definitions/SystemMessage"}, {"$ref": "#/definitions/FunctionMessage"}, {"$ref": "#/definitions/ToolMessage"}]}, "definitions": {"ToolCall": {"title": "ToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "object"}, "id": {"title": "Id", "type": "string"}, "type": {"title": "Type", "enum": ["tool_call"], "type": "string"}}, "required": ["name", "args", "id"]}, "InvalidToolCall": {"title": "InvalidToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "string"}, "id": {"title": "Id", "type": "string"}, "error": {"title": "Error", "type": "string"}, "type": {"title": "Type", "enum": ["invalid_tool_call"], "type": "string"}}, "required": ["name", "args", "id", "error"]}, "UsageMetadata": {"title": "UsageMetadata", "type": "object", "properties": {"input_tokens": {"title": "Input Tokens", "type": "integer"}, "output_tokens": {"title": "Output Tokens", "type": "integer"}, "total_tokens": {"title": "Total Tokens", "type": "integer"}}, "required": ["input_tokens", "output_tokens", "total_tokens"]}, "AIMessage": {"title": "AIMessage", "description": "Message from an AI.\\n\\nAIMessage is returned from a chat model as a response to a prompt.\\n\\nThis message represents the output of the model and consists of both\\nthe raw output as returned by the model together standardized fields\\n(e.g., tool calls, usage metadata) added by the LangChain framework.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "ai", "enum": ["ai"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}, "tool_calls": {"title": "Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/ToolCall"}}, "invalid_tool_calls": {"title": "Invalid Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/InvalidToolCall"}}, "usage_metadata": {"$ref": "#/definitions/UsageMetadata"}}, "required": ["content"]}, "HumanMessage": {"title": "HumanMessage", "description": "Message from a human.\\n\\nHumanMessages are messages that are passed in from a human to the model.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import HumanMessage, SystemMessage\\n\\n messages = [\\n SystemMessage(\\n content=\\"You are a helpful assistant! Your name is Bob.\\"\\n ),\\n HumanMessage(\\n content=\\"What is your name?\\"\\n )\\n ]\\n\\n # Instantiate a chat model and invoke it with the messages\\n model = ...\\n print(model.invoke(messages))", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "human", "enum": ["human"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}}, "required": ["content"]}, "ChatMessage": {"title": "ChatMessage", "description": "Message that can be assigned an arbitrary speaker (i.e. role).", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "chat", "enum": ["chat"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "role": {"title": "Role", "type": "string"}}, "required": ["content", "role"]}, "SystemMessage": {"title": "SystemMessage", "description": "Message for priming AI behavior.\\n\\nThe system message is usually passed in as the first of a sequence\\nof input messages.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import HumanMessage, SystemMessage\\n\\n messages = [\\n SystemMessage(\\n content=\\"You are a helpful assistant! Your name is Bob.\\"\\n ),\\n HumanMessage(\\n content=\\"What is your name?\\"\\n )\\n ]\\n\\n # Define a chat model and invoke it with the messages\\n print(model.invoke(messages))", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "system", "enum": ["system"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content"]}, "FunctionMessage": {"title": "FunctionMessage", "description": "Message for passing the result of executing a tool back to a model.\\n\\nFunctionMessage are an older version of the ToolMessage schema, and\\ndo not contain the tool_call_id field.\\n\\nThe tool_call_id field is used to associate the tool call request with the\\ntool call response. This is useful in situations where a chat model is able\\nto request multiple tool calls in parallel.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "function", "enum": ["function"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content", "name"]}, "ToolMessage": {"title": "ToolMessage", "description": "Message for passing the result of executing a tool back to a model.\\n\\nToolMessages contain the result of a tool invocation. Typically, the result\\nis encoded inside the `content` field.\\n\\nExample: A ToolMessage representing a result of 42 from a tool call with id\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import ToolMessage\\n\\n ToolMessage(content=\'42\', tool_call_id=\'call_Jja7J89XsjrOLA5r!MEOW!SL\')\\n\\n\\nExample: A ToolMessage where only part of the tool output is sent to the model\\n and the full output is passed in to artifact.\\n\\n .. versionadded:: 0.2.17\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import ToolMessage\\n\\n tool_output = {\\n \\"stdout\\": \\"From the graph we can see that the correlation between x and y is ...\\",\\n \\"stderr\\": None,\\n \\"artifacts\\": {\\"type\\": \\"image\\", \\"base64_data\\": \\"/9j/4gIcSU...\\"},\\n }\\n\\n ToolMessage(\\n content=tool_output[\\"stdout\\"],\\n artifact=tool_output,\\n tool_call_id=\'call_Jja7J89XsjrOLA5r!MEOW!SL\',\\n )\\n\\nThe tool_call_id field is used to associate the tool call request with the\\ntool call response. This is useful in situations where a chat model is able\\nto request multiple tool calls in parallel.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "tool", "enum": ["tool"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "tool_call_id": {"title": "Tool Call Id", "type": "string"}, "artifact": {"title": "Artifact"}, "status": {"title": "Status", "default": "success", "enum": ["success", "error"], "type": "string"}}, "required": ["content", "tool_call_id"]}}}' +# --- +# name: test_message_graph[postgres_pool].1 + '{"title": "LangGraphOutput", "type": "array", "items": {"anyOf": [{"$ref": "#/definitions/AIMessage"}, {"$ref": "#/definitions/HumanMessage"}, {"$ref": "#/definitions/ChatMessage"}, {"$ref": "#/definitions/SystemMessage"}, {"$ref": "#/definitions/FunctionMessage"}, {"$ref": "#/definitions/ToolMessage"}]}, "definitions": {"ToolCall": {"title": "ToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "object"}, "id": {"title": "Id", "type": "string"}, "type": {"title": "Type", "enum": ["tool_call"], "type": "string"}}, "required": ["name", "args", "id"]}, "InvalidToolCall": {"title": "InvalidToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "string"}, "id": {"title": "Id", "type": "string"}, "error": {"title": "Error", "type": "string"}, "type": {"title": "Type", "enum": ["invalid_tool_call"], "type": "string"}}, "required": ["name", "args", "id", "error"]}, "UsageMetadata": {"title": "UsageMetadata", "type": "object", "properties": {"input_tokens": {"title": "Input Tokens", "type": "integer"}, "output_tokens": {"title": "Output Tokens", "type": "integer"}, "total_tokens": {"title": "Total Tokens", "type": "integer"}}, "required": ["input_tokens", "output_tokens", "total_tokens"]}, "AIMessage": {"title": "AIMessage", "description": "Message from an AI.\\n\\nAIMessage is returned from a chat model as a response to a prompt.\\n\\nThis message represents the output of the model and consists of both\\nthe raw output as returned by the model together standardized fields\\n(e.g., tool calls, usage metadata) added by the LangChain framework.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "ai", "enum": ["ai"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}, "tool_calls": {"title": "Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/ToolCall"}}, "invalid_tool_calls": {"title": "Invalid Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/InvalidToolCall"}}, "usage_metadata": {"$ref": "#/definitions/UsageMetadata"}}, "required": ["content"]}, "HumanMessage": {"title": "HumanMessage", "description": "Message from a human.\\n\\nHumanMessages are messages that are passed in from a human to the model.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import HumanMessage, SystemMessage\\n\\n messages = [\\n SystemMessage(\\n content=\\"You are a helpful assistant! Your name is Bob.\\"\\n ),\\n HumanMessage(\\n content=\\"What is your name?\\"\\n )\\n ]\\n\\n # Instantiate a chat model and invoke it with the messages\\n model = ...\\n print(model.invoke(messages))", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "human", "enum": ["human"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}}, "required": ["content"]}, "ChatMessage": {"title": "ChatMessage", "description": "Message that can be assigned an arbitrary speaker (i.e. role).", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "chat", "enum": ["chat"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "role": {"title": "Role", "type": "string"}}, "required": ["content", "role"]}, "SystemMessage": {"title": "SystemMessage", "description": "Message for priming AI behavior.\\n\\nThe system message is usually passed in as the first of a sequence\\nof input messages.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import HumanMessage, SystemMessage\\n\\n messages = [\\n SystemMessage(\\n content=\\"You are a helpful assistant! Your name is Bob.\\"\\n ),\\n HumanMessage(\\n content=\\"What is your name?\\"\\n )\\n ]\\n\\n # Define a chat model and invoke it with the messages\\n print(model.invoke(messages))", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "system", "enum": ["system"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content"]}, "FunctionMessage": {"title": "FunctionMessage", "description": "Message for passing the result of executing a tool back to a model.\\n\\nFunctionMessage are an older version of the ToolMessage schema, and\\ndo not contain the tool_call_id field.\\n\\nThe tool_call_id field is used to associate the tool call request with the\\ntool call response. This is useful in situations where a chat model is able\\nto request multiple tool calls in parallel.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "function", "enum": ["function"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content", "name"]}, "ToolMessage": {"title": "ToolMessage", "description": "Message for passing the result of executing a tool back to a model.\\n\\nToolMessages contain the result of a tool invocation. Typically, the result\\nis encoded inside the `content` field.\\n\\nExample: A ToolMessage representing a result of 42 from a tool call with id\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import ToolMessage\\n\\n ToolMessage(content=\'42\', tool_call_id=\'call_Jja7J89XsjrOLA5r!MEOW!SL\')\\n\\n\\nExample: A ToolMessage where only part of the tool output is sent to the model\\n and the full output is passed in to artifact.\\n\\n .. versionadded:: 0.2.17\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import ToolMessage\\n\\n tool_output = {\\n \\"stdout\\": \\"From the graph we can see that the correlation between x and y is ...\\",\\n \\"stderr\\": None,\\n \\"artifacts\\": {\\"type\\": \\"image\\", \\"base64_data\\": \\"/9j/4gIcSU...\\"},\\n }\\n\\n ToolMessage(\\n content=tool_output[\\"stdout\\"],\\n artifact=tool_output,\\n tool_call_id=\'call_Jja7J89XsjrOLA5r!MEOW!SL\',\\n )\\n\\nThe tool_call_id field is used to associate the tool call request with the\\ntool call response. This is useful in situations where a chat model is able\\nto request multiple tool calls in parallel.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "tool", "enum": ["tool"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "tool_call_id": {"title": "Tool Call Id", "type": "string"}, "artifact": {"title": "Artifact"}, "status": {"title": "Status", "default": "success", "enum": ["success", "error"], "type": "string"}}, "required": ["content", "tool_call_id"]}}}' +# --- +# name: test_message_graph[postgres_pool].2 + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": "__start__" + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": [ + "tests", + "test_pregel", + "FakeFuntionChatModel" + ], + "name": "agent" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "prebuilt", + "tool_node", + "ToolNode" + ], + "name": "tools" + } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" + } + ], + "edges": [ + { + "source": "__start__", + "target": "agent" + }, + { + "source": "tools", + "target": "agent" + }, + { + "source": "agent", + "target": "tools", + "data": "continue", + "conditional": true + }, + { + "source": "agent", + "target": "__end__", + "data": "end", + "conditional": true + } + ] + } + ''' +# --- +# name: test_message_graph[postgres_pool].3 + ''' + graph TD; + __start__ --> agent; + tools --> agent; + agent -.  continue  .-> tools; + agent -.  end  .-> __end__; + + ''' +# --- # name: test_message_graph[sqlite] '{"title": "LangGraphInput", "type": "array", "items": {"anyOf": [{"$ref": "#/definitions/AIMessage"}, {"$ref": "#/definitions/HumanMessage"}, {"$ref": "#/definitions/ChatMessage"}, {"$ref": "#/definitions/SystemMessage"}, {"$ref": "#/definitions/FunctionMessage"}, {"$ref": "#/definitions/ToolMessage"}]}, "definitions": {"ToolCall": {"title": "ToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "object"}, "id": {"title": "Id", "type": "string"}, "type": {"title": "Type", "enum": ["tool_call"], "type": "string"}}, "required": ["name", "args", "id"]}, "InvalidToolCall": {"title": "InvalidToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "string"}, "id": {"title": "Id", "type": "string"}, "error": {"title": "Error", "type": "string"}, "type": {"title": "Type", "enum": ["invalid_tool_call"], "type": "string"}}, "required": ["name", "args", "id", "error"]}, "UsageMetadata": {"title": "UsageMetadata", "type": "object", "properties": {"input_tokens": {"title": "Input Tokens", "type": "integer"}, "output_tokens": {"title": "Output Tokens", "type": "integer"}, "total_tokens": {"title": "Total Tokens", "type": "integer"}}, "required": ["input_tokens", "output_tokens", "total_tokens"]}, "AIMessage": {"title": "AIMessage", "description": "Message from an AI.\\n\\nAIMessage is returned from a chat model as a response to a prompt.\\n\\nThis message represents the output of the model and consists of both\\nthe raw output as returned by the model together standardized fields\\n(e.g., tool calls, usage metadata) added by the LangChain framework.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "ai", "enum": ["ai"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}, "tool_calls": {"title": "Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/ToolCall"}}, "invalid_tool_calls": {"title": "Invalid Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/InvalidToolCall"}}, "usage_metadata": {"$ref": "#/definitions/UsageMetadata"}}, "required": ["content"]}, "HumanMessage": {"title": "HumanMessage", "description": "Message from a human.\\n\\nHumanMessages are messages that are passed in from a human to the model.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import HumanMessage, SystemMessage\\n\\n messages = [\\n SystemMessage(\\n content=\\"You are a helpful assistant! Your name is Bob.\\"\\n ),\\n HumanMessage(\\n content=\\"What is your name?\\"\\n )\\n ]\\n\\n # Instantiate a chat model and invoke it with the messages\\n model = ...\\n print(model.invoke(messages))", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "human", "enum": ["human"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}}, "required": ["content"]}, "ChatMessage": {"title": "ChatMessage", "description": "Message that can be assigned an arbitrary speaker (i.e. role).", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "chat", "enum": ["chat"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "role": {"title": "Role", "type": "string"}}, "required": ["content", "role"]}, "SystemMessage": {"title": "SystemMessage", "description": "Message for priming AI behavior.\\n\\nThe system message is usually passed in as the first of a sequence\\nof input messages.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import HumanMessage, SystemMessage\\n\\n messages = [\\n SystemMessage(\\n content=\\"You are a helpful assistant! Your name is Bob.\\"\\n ),\\n HumanMessage(\\n content=\\"What is your name?\\"\\n )\\n ]\\n\\n # Define a chat model and invoke it with the messages\\n print(model.invoke(messages))", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "system", "enum": ["system"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content"]}, "FunctionMessage": {"title": "FunctionMessage", "description": "Message for passing the result of executing a tool back to a model.\\n\\nFunctionMessage are an older version of the ToolMessage schema, and\\ndo not contain the tool_call_id field.\\n\\nThe tool_call_id field is used to associate the tool call request with the\\ntool call response. This is useful in situations where a chat model is able\\nto request multiple tool calls in parallel.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "function", "enum": ["function"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content", "name"]}, "ToolMessage": {"title": "ToolMessage", "description": "Message for passing the result of executing a tool back to a model.\\n\\nToolMessages contain the result of a tool invocation. Typically, the result\\nis encoded inside the `content` field.\\n\\nExample: A ToolMessage representing a result of 42 from a tool call with id\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import ToolMessage\\n\\n ToolMessage(content=\'42\', tool_call_id=\'call_Jja7J89XsjrOLA5r!MEOW!SL\')\\n\\n\\nExample: A ToolMessage where only part of the tool output is sent to the model\\n and the full output is passed in to artifact.\\n\\n .. versionadded:: 0.2.17\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import ToolMessage\\n\\n tool_output = {\\n \\"stdout\\": \\"From the graph we can see that the correlation between x and y is ...\\",\\n \\"stderr\\": None,\\n \\"artifacts\\": {\\"type\\": \\"image\\", \\"base64_data\\": \\"/9j/4gIcSU...\\"},\\n }\\n\\n ToolMessage(\\n content=tool_output[\\"stdout\\"],\\n artifact=tool_output,\\n tool_call_id=\'call_Jja7J89XsjrOLA5r!MEOW!SL\',\\n )\\n\\nThe tool_call_id field is used to associate the tool call request with the\\ntool call response. This is useful in situations where a chat model is able\\nto request multiple tool calls in parallel.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "tool", "enum": ["tool"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "tool_call_id": {"title": "Tool Call Id", "type": "string"}, "artifact": {"title": "Artifact"}, "status": {"title": "Status", "default": "success", "enum": ["success", "error"], "type": "string"}}, "required": ["content", "tool_call_id"]}}}' # --- @@ -4299,6 +5006,24 @@ ''' # --- +# name: test_start_branch_then[postgres_pool] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__]):::first + tool_two_slow(tool_two_slow) + tool_two_fast(tool_two_fast) + __end__([__end__]):::last + __start__ -.-> tool_two_slow; + tool_two_slow --> __end__; + __start__ -.-> tool_two_fast; + tool_two_fast --> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- # name: test_start_branch_then[sqlite] ''' %%{init: {'flowchart': {'curve': 'linear'}}}%% diff --git a/libs/langgraph/tests/conftest.py b/libs/langgraph/tests/conftest.py index 1219ff0f1..bb6ef9603 100644 --- a/libs/langgraph/tests/conftest.py +++ b/libs/langgraph/tests/conftest.py @@ -7,6 +7,7 @@ from uuid import UUID, uuid4 import pytest from psycopg import AsyncConnection, Connection +from psycopg_pool import AsyncConnectionPool, ConnectionPool from pytest_mock import MockerFixture from langgraph.checkpoint.postgres import PostgresSaver @@ -122,6 +123,26 @@ def checkpointer_postgres_pipe(): conn.execute(f"DROP DATABASE {database}") +@pytest.fixture(scope="function") +def checkpointer_postgres_pool(): + database = f"test_{uuid4().hex[:16]}" + # create unique db + with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn: + conn.execute(f"CREATE DATABASE {database}") + try: + # yield checkpointer + with ConnectionPool( + DEFAULT_POSTGRES_URI + database, max_size=10, kwargs={"autocommit": True} + ) as pool: + checkpointer = PostgresSaver(pool) + checkpointer.setup() + yield checkpointer + finally: + # drop unique db + with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn: + conn.execute(f"DROP DATABASE {database}") + + @pytest.fixture(scope="function") def checkpointer_postgres_aio(): if sys.version_info < (3, 10): @@ -185,3 +206,51 @@ async def _checkpointer_postgres_aio_pipe(): DEFAULT_POSTGRES_URI, autocommit=True ) as conn: await conn.execute(f"DROP DATABASE {database}") + + +@pytest.fixture(scope="function") +def checkpointer_postgres_aio_pool(): + if sys.version_info < (3, 10): + pytest.skip("Async Postgres tests require Python 3.10+") + with agen_to_gen(_checkpointer_postgres_aio_pool()) as checkpointer: + yield checkpointer + + +@asynccontextmanager +async def _checkpointer_postgres_aio_pool(): + database = f"test_{uuid4().hex[:16]}" + # create unique db + async with await AsyncConnection.connect( + DEFAULT_POSTGRES_URI, autocommit=True + ) as conn: + await conn.execute(f"CREATE DATABASE {database}") + try: + # yield checkpointer + async with AsyncConnectionPool( + DEFAULT_POSTGRES_URI + database, max_size=10, kwargs={"autocommit": True} + ) as pool: + checkpointer = AsyncPostgresSaver(pool) + await checkpointer.setup() + yield checkpointer + finally: + # drop unique db + async with await AsyncConnection.connect( + DEFAULT_POSTGRES_URI, autocommit=True + ) as conn: + await conn.execute(f"DROP DATABASE {database}") + + +ALL_CHECKPOINTERS_SYNC = [ + "memory", + "sqlite", + "postgres", + "postgres_pipe", + "postgres_pool", +] +ALL_CHECKPOINTERS_ASYNC = [ + "memory", + "sqlite_aio", + "postgres_aio", + "postgres_aio_pipe", + "postgres_aio_pool", +] diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 685eb7c87..2c28c4ec0 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -70,6 +70,7 @@ from langgraph.pregel.retry import RetryPolicy from langgraph.pregel.types import PregelTask from langgraph.store.memory import MemoryStore from tests.any_str import AnyStr, AnyVersion, UnsortedSequence +from tests.conftest import ALL_CHECKPOINTERS_SYNC from tests.fake_tracer import FakeTracer from tests.memory_assert import ( MemorySaverAssertCheckpointMetadata, @@ -593,10 +594,7 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: assert step == 2 -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite", "postgres", "postgres_pipe"], -) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_invoke_two_processes_in_out_interrupt( request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture ) -> None: @@ -804,10 +802,7 @@ def test_invoke_two_processes_in_out_interrupt( ] -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite", "postgres", "postgres_pipe"], -) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_fork_always_re_runs_nodes( request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture ) -> None: @@ -1248,10 +1243,7 @@ def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) -> Non assert app.invoke(2) == [3, 3] -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite", "postgres", "postgres_pipe"], -) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_invoke_checkpoint_two( mocker: MockerFixture, request: pytest.FixtureRequest, checkpointer_name: str ) -> None: @@ -1324,10 +1316,7 @@ def test_invoke_checkpoint_two( assert checkpoint["channel_values"].get("total") == 5 -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite", "postgres", "postgres_pipe"], -) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_pending_writes_resume( request: pytest.FixtureRequest, checkpointer_name: str ) -> None: @@ -1606,10 +1595,7 @@ async def test_checkpointer_null_pending_writes() -> None: ] * 4 -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite", "postgres", "postgres_pipe"], -) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_invoke_checkpoint_three( mocker: MockerFixture, request: pytest.FixtureRequest, checkpointer_name: str ) -> None: @@ -1939,10 +1925,7 @@ def test_channel_enter_exit_timing(mocker: MockerFixture) -> None: assert cleanup.call_count == 1, "Expected cleanup to be called once" -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite", "postgres", "postgres_pipe"], -) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_conditional_graph( snapshot: SnapshotAssertion, request: pytest.FixtureRequest, checkpointer_name: str ) -> None: @@ -2802,10 +2785,7 @@ def test_conditional_entrypoint_to_multiple_state_graph( } -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite", "postgres", "postgres_pipe"], -) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_conditional_state_graph( snapshot: SnapshotAssertion, mocker: MockerFixture, @@ -4060,10 +4040,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: ] -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite", "postgres", "postgres_pipe"], -) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_state_graph_packets( request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture ) -> None: @@ -4592,10 +4569,7 @@ def test_state_graph_packets( ) -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite", "postgres", "postgres_pipe"], -) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_message_graph( snapshot: SnapshotAssertion, deterministic_uuids: MockerFixture, @@ -5323,10 +5297,7 @@ def test_message_graph( ) -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite", "postgres", "postgres_pipe"], -) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_root_graph( deterministic_uuids: MockerFixture, request: pytest.FixtureRequest, @@ -6400,10 +6371,7 @@ def test_in_one_fan_out_out_one_graph_state() -> None: ] -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite", "postgres", "postgres_pipe"], -) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_dynamic_interrupt( request: pytest.FixtureRequest, checkpointer_name: str ) -> None: @@ -6487,10 +6455,7 @@ def test_dynamic_interrupt( ) -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite", "postgres", "postgres_pipe"], -) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_start_branch_then( snapshot: SnapshotAssertion, request: pytest.FixtureRequest, checkpointer_name: str ) -> None: @@ -6680,10 +6645,7 @@ def test_start_branch_then( ) -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite", "postgres", "postgres_pipe"], -) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_branch_then( snapshot: SnapshotAssertion, request: pytest.FixtureRequest, checkpointer_name: str ) -> None: @@ -7210,10 +7172,7 @@ def test_branch_then( ) -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite", "postgres", "postgres_pipe"], -) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_in_one_fan_out_state_graph_waiting_edge( snapshot: SnapshotAssertion, request: pytest.FixtureRequest, checkpointer_name: str ) -> None: @@ -7341,10 +7300,7 @@ def test_in_one_fan_out_state_graph_waiting_edge( ] -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite", "postgres", "postgres_pipe"], -) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_in_one_fan_out_state_graph_waiting_edge_via_branch( snapshot: SnapshotAssertion, request: pytest.FixtureRequest, checkpointer_name: str ) -> None: @@ -7438,10 +7394,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_via_branch( ] -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite", "postgres", "postgres_pipe"], -) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1( snapshot: SnapshotAssertion, mocker: MockerFixture, @@ -7609,10 +7562,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1( } -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite", "postgres", "postgres_pipe"], -) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2( snapshot: SnapshotAssertion, mocker: MockerFixture, @@ -7779,10 +7729,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2( } -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite", "postgres", "postgres_pipe"], -) +@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 ) -> None: @@ -8254,10 +8201,7 @@ def test_nested_graph(snapshot: SnapshotAssertion) -> None: @pytest.mark.repeat(10) -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite", "postgres", "postgres_pipe"], -) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_nested_graph_interrupts( request: pytest.FixtureRequest, checkpointer_name: str ) -> None: @@ -9949,10 +9893,7 @@ def test_nested_graph_interrupts( ] -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite", "postgres", "postgres_pipe"], -) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_nested_graph_interrupts_parallel( request: pytest.FixtureRequest, checkpointer_name: str ) -> None: @@ -10076,10 +10017,7 @@ def test_nested_graph_interrupts_parallel( ] -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite", "postgres", "postgres_pipe"], -) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_doubly_nested_graph_interrupts( request: pytest.FixtureRequest, checkpointer_name: str ) -> None: @@ -10174,10 +10112,7 @@ def test_doubly_nested_graph_interrupts( ] -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite", "postgres", "postgres_pipe"], -) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_nested_graph_state( request: pytest.FixtureRequest, checkpointer_name: str ) -> None: @@ -10609,10 +10544,7 @@ def test_nested_graph_state( assert app.get_state(actual_snapshot.config) == expected_snapshot -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite", "postgres", "postgres_pipe"], -) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_doubly_nested_graph_state( request: pytest.FixtureRequest, checkpointer_name: str ) -> None: @@ -10867,11 +10799,7 @@ def test_doubly_nested_graph_state( ) -@pytest.mark.repeat(10) -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite", "postgres", "postgres_pipe"], -) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_send_to_nested_graphs( request: pytest.FixtureRequest, checkpointer_name: str ) -> None: @@ -11320,10 +11248,7 @@ def test_checkpoint_metadata() -> None: assert chkpnt_tuple.metadata["test_config_4"] == "bar" -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite", "postgres", "postgres_pipe"], -) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_remove_message_via_state_update( request: pytest.FixtureRequest, checkpointer_name: str ) -> None: @@ -11507,10 +11432,7 @@ def test_xray_lance(snapshot: SnapshotAssertion): assert graph.get_graph(xray=1).to_json() == snapshot -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite", "postgres", "postgres_pipe"], -) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_channel_values(request: pytest.FixtureRequest, checkpointer_name: str) -> None: checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 50aafae86..06d9578f3 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -70,6 +70,7 @@ from langgraph.pregel.retry import RetryPolicy from langgraph.pregel.types import PregelTask from langgraph.store.memory import MemoryStore from tests.any_str import AnyStr, AnyVersion, UnsortedSequence +from tests.conftest import ALL_CHECKPOINTERS_ASYNC from tests.fake_tracer import FakeTracer from tests.memory_assert import ( MemorySaverAssertCheckpointMetadata, @@ -214,10 +215,7 @@ async def test_node_cancellation_on_other_node_exception() -> None: assert inner_task_cancelled -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"], -) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_dynamic_interrupt( checkpointer_name: str, snapshot: SnapshotAssertion, request: pytest.FixtureRequest ) -> None: @@ -305,10 +303,7 @@ async def test_dynamic_interrupt( # TODO use aget_state_history -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"], -) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_node_not_cancelled_on_other_node_interrupted( checkpointer_name: str, request: pytest.FixtureRequest ) -> None: @@ -390,10 +385,7 @@ async def test_step_timeout_on_stream_hang() -> None: assert inner_task_cancelled -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"], -) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_cancel_graph_astream( request: pytest.FixtureRequest, checkpointer_name: str ) -> None: @@ -463,10 +455,7 @@ async def test_cancel_graph_astream( assert state.metadata == {"source": "loop", "step": 0, "writes": None} -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe", None], -) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_cancel_graph_astream_events_v2( request: pytest.FixtureRequest, checkpointer_name: Optional[str] ) -> None: @@ -820,10 +809,7 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: assert step == 2 -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"], -) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_invoke_two_processes_in_out_interrupt( request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture ) -> None: @@ -1038,10 +1024,7 @@ async def test_invoke_two_processes_in_out_interrupt( ] -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"], -) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_fork_always_re_runs_nodes( request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture ) -> None: @@ -1546,10 +1529,7 @@ async def test_invoke_checkpoint(mocker: MockerFixture) -> None: assert checkpoint["channel_values"].get("total") == 5 -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"], -) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_pending_writes_resume( request: pytest.FixtureRequest, checkpointer_name: str ) -> None: @@ -1810,10 +1790,7 @@ async def test_cond_edge_after_send() -> None: assert await graph.ainvoke(["0"]) == ["0", "1", "2", "2", "3"] -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"], -) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_invoke_checkpoint_three( mocker: MockerFixture, request: pytest.FixtureRequest, checkpointer_name: str ) -> None: @@ -2155,10 +2132,7 @@ async def test_channel_enter_exit_timing(mocker: MockerFixture) -> None: assert cleanup_async.call_count == 1, "Expected cleanup to be called once" -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"], -) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_conditional_graph( request: pytest.FixtureRequest, checkpointer_name: str ) -> None: @@ -5036,10 +5010,7 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: ] -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"], -) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_start_branch_then( request: pytest.FixtureRequest, checkpointer_name: str ) -> None: @@ -5243,10 +5214,7 @@ async def test_start_branch_then( ) -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"], -) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_branch_then( request: pytest.FixtureRequest, checkpointer_name: str ) -> None: @@ -6730,10 +6698,7 @@ async def test_nested_graph(snapshot: SnapshotAssertion) -> None: @pytest.mark.repeat(10) -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"], -) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_nested_graph_interrupts( request: pytest.FixtureRequest, checkpointer_name: str ) -> None: @@ -8443,10 +8408,7 @@ async def test_nested_graph_interrupts( ] -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"], -) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_nested_graph_interrupts_parallel( request: pytest.FixtureRequest, checkpointer_name: str ) -> None: @@ -8572,10 +8534,7 @@ async def test_nested_graph_interrupts_parallel( ] -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"], -) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_doubly_nested_graph_interrupts( request: pytest.FixtureRequest, checkpointer_name: str ) -> None: @@ -8673,10 +8632,7 @@ async def test_doubly_nested_graph_interrupts( ] -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"], -) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_nested_graph_state( request: pytest.FixtureRequest, checkpointer_name: str ) -> None: @@ -9113,10 +9069,7 @@ async def test_nested_graph_state( assert await app.aget_state(actual_snapshot.config) == expected_snapshot -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"], -) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_doubly_nested_graph_state( request: pytest.FixtureRequest, checkpointer_name: str ) -> None: @@ -9375,11 +9328,7 @@ async def test_doubly_nested_graph_state( ) -@pytest.mark.repeat(10) -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"], -) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_send_to_nested_graphs( request: pytest.FixtureRequest, checkpointer_name: str ) -> None: