diff --git a/docs/mkdocs.yml b/docs/mkdocs.yml index 1503e4a2d..22d1b7e67 100644 --- a/docs/mkdocs.yml +++ b/docs/mkdocs.yml @@ -380,6 +380,7 @@ nav: - Types: reference/types.md - Constants: reference/constants.md - Pregel: reference/pregel.md + - Config: reference/config.md - Functional API: reference/func.md - LangGraph Platform: - Server API: "cloud/reference/api/api_ref.md" diff --git a/libs/langgraph/langgraph/config.py b/libs/langgraph/langgraph/config.py new file mode 100644 index 000000000..df2b637b7 --- /dev/null +++ b/libs/langgraph/langgraph/config.py @@ -0,0 +1,180 @@ +import asyncio +import sys + +from langchain_core.runnables import RunnableConfig +from langchain_core.runnables.config import var_child_runnable_config + +from langgraph.constants import CONF, CONFIG_KEY_STORE, CONFIG_KEY_STREAM_WRITER +from langgraph.store.base import BaseStore +from langgraph.types import StreamWriter + + +def get_config() -> RunnableConfig: + if sys.version_info < (3, 11): + try: + if asyncio.current_task(): + raise RuntimeError( + "Python 3.11 or later required to use this in an async context" + ) + except RuntimeError: + pass + if var_config := var_child_runnable_config.get(): + return var_config + else: + raise RuntimeError("Called get_config outside of a runnable context") + + +def get_store() -> BaseStore: + """Access LangGraph store from inside a graph node or entrypoint task at runtime. + + Can be called from inside any [StateGraph][langgraph.graph.StateGraph] node or + functional API [task][langgraph.func.task], as long as the StateGraph or the [entrypoint][langgraph.func.entrypoint] + was initialized with a store, e.g.: + + ```python + # with StateGraph + graph = ( + StateGraph(...) + ... + .compile(store=store) + ) + + # or with entrypoint + @entrypoint(store=store) + def workflow(inputs): + ... + ``` + + !!! warning "Async with Python < 3.11" + + If you are using Python < 3.11 and are running LangGraph asynchronously, + `get_store()` won't work since it uses [contextvar](https://docs.python.org/3/library/contextvars.html) propagation (only available in [Python >= 3.11](https://docs.python.org/3/library/asyncio-task.html#asyncio.create_task)). + + + Example: Using with StateGraph + ```python + from typing_extensions import TypedDict + from langgraph.graph import StateGraph, START + from langgraph.store.memory import InMemoryStore + from langgraph.config import get_store + + store = InMemoryStore() + store.put(("values",), "foo", {"bar": 2}) + + class State(TypedDict): + foo: int + + def my_node(state: State): + my_store = get_store() + stored_value = my_store.get(("values",), "foo").value["bar"] + return {"foo": stored_value + 1} + + graph = ( + StateGraph(State) + .add_node(my_node) + .add_edge(START, "my_node") + .compile(store=store) + ) + + graph.invoke({"foo": 1}) + ``` + + ```pycon + {'foo': 3} + ``` + + Example: Using with functional API + ```python + from langgraph.func import entrypoint, task + from langgraph.store.memory import InMemoryStore + from langgraph.config import get_store + + store = InMemoryStore() + store.put(("values",), "foo", {"bar": 2}) + + @task + def my_task(value: int): + my_store = get_store() + stored_value = my_store.get(("values",), "foo").value["bar"] + return stored_value + 1 + + @entrypoint(store=store) + def workflow(value: int): + return my_task(value).result() + + workflow.invoke(1) + ``` + + ```pycon + 3 + ``` + """ + config = get_config() + return config[CONF][CONFIG_KEY_STORE] + + +def get_stream_writer() -> StreamWriter: + """Access LangGraph [StreamWriter][langgraph.types.StreamWriter] from inside a graph node or entrypoint task at runtime. + + Can be called from inside any [StateGraph][langgraph.graph.StateGraph] node or + functional API [task][langgraph.func.task]. + + !!! warning "Async with Python < 3.11" + + If you are using Python < 3.11 and are running LangGraph asynchronously, + `get_stream_writer()` won't work since it uses [contextvar](https://docs.python.org/3/library/contextvars.html) propagation (only available in [Python >= 3.11](https://docs.python.org/3/library/asyncio-task.html#asyncio.create_task)). + + Example: Using with StateGraph + ```python + from typing_extensions import TypedDict + from langgraph.graph import StateGraph, START + from langgraph.config import get_stream_writer + + class State(TypedDict): + foo: int + + def my_node(state: State): + my_stream_writer = get_stream_writer() + my_stream_writer({"custom_data": "Hello!"}) + return {"foo": state["foo"] + 1} + + graph = ( + StateGraph(State) + .add_node(my_node) + .add_edge(START, "my_node") + .compile(store=store) + ) + + for chunk in graph.stream({"foo": 1}, stream_mode="custom"): + print(chunk) + ``` + + ```pycon + {'custom_data': 'Hello!'} + ``` + + Example: Using with functional API + ```python + from langgraph.func import entrypoint, task + from langgraph.config import get_stream_writer + + @task + def my_task(value: int): + my_stream_writer = get_stream_writer() + my_stream_writer({"custom_data": "Hello!"}) + return value + 1 + + @entrypoint(store=store) + def workflow(value: int): + return my_task(value).result() + + for chunk in workflow.stream(1, stream_mode="custom"): + print(chunk) + ``` + + ```pycon + {'custom_data': 'Hello!'} + ``` + """ + config = get_config() + return config[CONF][CONFIG_KEY_STREAM_WRITER] diff --git a/libs/langgraph/langgraph/utils/config.py b/libs/langgraph/langgraph/utils/config.py index 55c89b3ba..baf1bb8a9 100644 --- a/libs/langgraph/langgraph/utils/config.py +++ b/libs/langgraph/langgraph/utils/config.py @@ -1,5 +1,3 @@ -import asyncio -import sys from collections import ChainMap from typing import Any, Optional, Sequence, cast @@ -18,16 +16,15 @@ from langchain_core.runnables.config import ( ) from langgraph.checkpoint.base import CheckpointMetadata +from langgraph.config import get_config, get_store, get_stream_writer # noqa from langgraph.constants import ( CONF, CONFIG_KEY_CHECKPOINT_ID, CONFIG_KEY_CHECKPOINT_MAP, CONFIG_KEY_CHECKPOINT_NS, - CONFIG_KEY_STORE, NS_END, NS_SEP, ) -from langgraph.store.base import BaseStore def recast_checkpoint_ns(ns: str) -> str: @@ -320,23 +317,3 @@ def ensure_config(*configs: Optional[RunnableConfig]) -> RunnableConfig: ): empty["metadata"][key] = value return empty - - -def get_config() -> RunnableConfig: - if sys.version_info < (3, 11): - try: - if asyncio.current_task(): - raise RuntimeError( - "Python 3.11 or later required to use this in an async context" - ) - except RuntimeError: - pass - if var_config := var_child_runnable_config.get(): - return var_config - else: - raise RuntimeError("Called get_config outside of a runnable context") - - -def get_store() -> BaseStore: - config = get_config() - return config[CONF][CONFIG_KEY_STORE] diff --git a/libs/langgraph/tests/test_prebuilt.py b/libs/langgraph/tests/test_prebuilt.py index 6a5327ca6..69751dc1a 100644 --- a/libs/langgraph/tests/test_prebuilt.py +++ b/libs/langgraph/tests/test_prebuilt.py @@ -62,6 +62,7 @@ from langgraph.prebuilt.tool_node import ( from langgraph.store.base import BaseStore from langgraph.store.memory import InMemoryStore from langgraph.types import Command, Interrupt, interrupt +from langgraph.utils.config import get_stream_writer from tests.conftest import ( ALL_CHECKPOINTERS_ASYNC, ALL_CHECKPOINTERS_SYNC, @@ -2214,3 +2215,57 @@ def test_react_with_subgraph_tools() -> None: content="What's 2 + 3 and 2 * 3?-What's 2 + 3 and 2 * 3?-5-6", id="1" ), ] + + +def test_tool_node_stream_writer() -> None: + @dec_tool + def streaming_tool(x: int) -> str: + """Do something with writer.""" + my_writer = get_stream_writer() + for value in ["foo", "bar", "baz"]: + my_writer({"custom_tool_value": value}) + + return x + + tool_node = ToolNode([streaming_tool]) + graph = ( + StateGraph(MessagesState) + .add_node("tools", tool_node) + .add_edge(START, "tools") + .compile() + ) + + tool_call = { + "name": "streaming_tool", + "args": {"x": 1}, + "id": "1", + "type": "tool_call", + } + inputs = { + "messages": [AIMessage("", tool_calls=[tool_call])], + } + + assert list(graph.stream(inputs, stream_mode="custom")) == [ + {"custom_tool_value": "foo"}, + {"custom_tool_value": "bar"}, + {"custom_tool_value": "baz"}, + ] + assert list(graph.stream(inputs, stream_mode=["custom", "updates"])) == [ + ("custom", {"custom_tool_value": "foo"}), + ("custom", {"custom_tool_value": "bar"}), + ("custom", {"custom_tool_value": "baz"}), + ( + "updates", + { + "tools": { + "messages": [ + _AnyIdToolMessage( + content="1", + name="streaming_tool", + tool_call_id="1", + ), + ], + }, + }, + ), + ]