codekingpro/portable-devtools

1import asyncio2import sys3from typing import Any4 5from langchain_core.runnables import RunnableConfig6from langchain_core.runnables.config import var_child_runnable_config7from langgraph.store.base import BaseStore8 9from langgraph._internal._constants import CONF, CONFIG_KEY_RUNTIME10from langgraph.types import StreamWriter11 12 13def _no_op_stream_writer(c: Any) -> None:14    pass15 16 17def get_config() -> RunnableConfig:18    if sys.version_info < (3, 11):19        try:20            if asyncio.current_task():21                raise RuntimeError(22                    "Python 3.11 or later required to use this in an async context"23                )24        except RuntimeError:25            pass26    if var_config := var_child_runnable_config.get():27        return var_config28    else:29        raise RuntimeError("Called get_config outside of a runnable context")30 31 32def get_store() -> BaseStore:33    """Access LangGraph store from inside a graph node or entrypoint task at runtime.34 35    Can be called from inside any [`StateGraph`][langgraph.graph.StateGraph] node or36    functional API [`task`][langgraph.func.task], as long as the `StateGraph` or the [`entrypoint`][langgraph.func.entrypoint]37    was initialized with a store, e.g.:38 39    ```python40    # with StateGraph41    graph = (42        StateGraph(...)43        ...44        .compile(store=store)45    )46 47    # or with entrypoint48    @entrypoint(store=store)49    def workflow(inputs):50        ...51    ```52 53    !!! warning "Async with Python < 3.11"54 55        If you are using Python < 3.11 and are running LangGraph asynchronously,56        `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)).57 58 59    Example: Using with `StateGraph`60        ```python61        from typing_extensions import TypedDict62        from langgraph.graph import StateGraph, START63        from langgraph.store.memory import InMemoryStore64        from langgraph.config import get_store65 66        store = InMemoryStore()67        store.put(("values",), "foo", {"bar": 2})68 69 70        class State(TypedDict):71            foo: int72 73 74        def my_node(state: State):75            my_store = get_store()76            stored_value = my_store.get(("values",), "foo").value["bar"]77            return {"foo": stored_value + 1}78 79 80        graph = (81            StateGraph(State)82            .add_node(my_node)83            .add_edge(START, "my_node")84            .compile(store=store)85        )86 87        graph.invoke({"foo": 1})88        ```89 90        ```pycon91        {"foo": 3}92        ```93 94    Example: Using with functional API95        ```python96        from langgraph.func import entrypoint, task97        from langgraph.store.memory import InMemoryStore98        from langgraph.config import get_store99 100        store = InMemoryStore()101        store.put(("values",), "foo", {"bar": 2})102 103 104        @task105        def my_task(value: int):106            my_store = get_store()107            stored_value = my_store.get(("values",), "foo").value["bar"]108            return stored_value + 1109 110 111        @entrypoint(store=store)112        def workflow(value: int):113            return my_task(value).result()114 115 116        workflow.invoke(1)117        ```118 119        ```pycon120        3121        ```122    """123    return get_config()[CONF][CONFIG_KEY_RUNTIME].store124 125 126def get_stream_writer() -> StreamWriter:127    """Access LangGraph [`StreamWriter`][langgraph.types.StreamWriter] from inside a graph node or entrypoint task at runtime.128 129    Can be called from inside any [`StateGraph`][langgraph.graph.StateGraph] node or130    functional API [`task`][langgraph.func.task].131 132    !!! warning "Async with Python < 3.11"133 134        If you are using Python < 3.11 and are running LangGraph asynchronously,135        `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)).136 137    Example: Using with `StateGraph`138        ```python139        from typing_extensions import TypedDict140        from langgraph.graph import StateGraph, START141        from langgraph.config import get_stream_writer142 143 144        class State(TypedDict):145            foo: int146 147 148        def my_node(state: State):149            my_stream_writer = get_stream_writer()150            my_stream_writer({"custom_data": "Hello!"})151            return {"foo": state["foo"] + 1}152 153 154        graph = (155            StateGraph(State)156            .add_node(my_node)157            .add_edge(START, "my_node")158            .compile(store=store)159        )160 161        for chunk in graph.stream({"foo": 1}, stream_mode="custom"):162            print(chunk)163        ```164 165        ```pycon166        {"custom_data": "Hello!"}167        ```168 169    Example: Using with functional API170        ```python171        from langgraph.func import entrypoint, task172        from langgraph.config import get_stream_writer173 174 175        @task176        def my_task(value: int):177            my_stream_writer = get_stream_writer()178            my_stream_writer({"custom_data": "Hello!"})179            return value + 1180 181 182        @entrypoint(store=store)183        def workflow(value: int):184            return my_task(value).result()185 186 187        for chunk in workflow.stream(1, stream_mode="custom"):188            print(chunk)189        ```190 191        ```pycon192        {"custom_data": "Hello!"}193        ```194    """195    runtime = get_config()[CONF][CONFIG_KEY_RUNTIME]196    return runtime.stream_writer197