WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-22 15:42:25 +02:00
Code Issues Packages Projects Releases Wiki Activity

langgraph: add get_stream_writer() (#3251)

Browse Source 
Alternative to #3250
This commit is contained in:
Vadym Barda
2025-01-30 20:45:44 +00:00
committed by GitHub
parent 80cef60405
commit a0ec9017f2
4 changed files with 237 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/mkdocs.yml
+1
View File
@@ -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"
libs/langgraph/langgraph/config.py
+180
View File
@@ -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]
libs/langgraph/langgraph/utils/config.py
+1 -24
View File
@@ -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]
libs/langgraph/tests/test_prebuilt.py
+55
View File
@@ -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",
),
],
},
},
),
]