Set store & stream writer #4044

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed

hinthornw wants to merge 1 commit into main from wfh/set_store

Closed

libs/langgraph/langgraph/config.py

-Original file line number
+Diff line change
@@ -1,4 +1,4 @@
     import asyncio
     import sys
     from typing import Any
@@ Expand Down Expand Up / @@ -118,6 +118,17 @@ @@
         return config[CONF][CONFIG_KEY_STORE]
+    def set_store(store: BaseStore) -> None:
+        """Set LangGraph store in context."""
+        var_config = var_child_runnable_config.get()
+        if not var_config:
+            var_config = {CONF: {CONFIG_KEY_STORE: store}}
+            var_child_runnable_config.set(var_config)
+        else:
+            var_config[CONF][CONFIG_KEY_STORE] = store
+            var_child_runnable_config.set(var_config)
     def get_stream_writer() -> StreamWriter:
         """Access LangGraph [StreamWriter][langgraph.types.StreamWriter] from inside a graph node or entrypoint task at runtime.
@@ Expand Down Expand Up / @@ -183,3 +194,14 @@ @@
         """
         config = get_config()
         return config[CONF].get(CONFIG_KEY_STREAM_WRITER, _no_op_stream_writer)
+    def set_stream_writer(stream_writer: StreamWriter) -> None:
+        """Set LangGraph stream writer in context."""
+        var_config = var_child_runnable_config.get()
+        if not var_config:
+            var_config = {CONF: {CONFIG_KEY_STREAM_WRITER: stream_writer}}
+            var_child_runnable_config.set(var_config)
+        else:
+            var_config[CONF][CONFIG_KEY_STREAM_WRITER] = stream_writer
+            var_child_runnable_config.set(var_config)

Provide feedback