Feat/support lg 06 context

typescript-sdk/integrations/langgraph/examples/python/pyproject.toml

@@ -21,7 +21,7 @@ langchain-experimental = ">=0.0.11"
 langchain-google-genai = ">=2.1.9"
 langchain-openai = ">=0.0.1"
 langgraph = "^0.6.1"
-ag-ui-langgraph = { version = "0.0.10", extras = ["fastapi"] }
+ag-ui-langgraph = { version = "0.0.12a1", extras = ["fastapi"] }
 python-dotenv = "^1.0.0"
 fastapi = "^0.115.12"
 

typescript-sdk/integrations/langgraph/python/ag_ui_langgraph/agent.py

@@ -1,6 +1,7 @@
 import uuid
 import json
-from typing import Optional, List, Any, Union, AsyncGenerator, Generator
+from typing import Optional, List, Any, Union, AsyncGenerator, Generator, Literal, Dict
+import inspect
 
 from langgraph.graph.state import CompiledStateGraph
 from langchain.schema import BaseMessage, SystemMessage
@@ -335,13 +336,15 @@ async def prepare_stream(self, input: RunAgentInput, agent_state: State, config:
 
         subgraphs_stream_enabled = input.forwarded_props.get('stream_subgraphs') if input.forwarded_props else False
 
-        stream = self.graph.astream_events(
-            stream_input,
+        kwargs = self.get_stream_kwargs(
+            input=stream_input,
             config=config,
-            subgraps=bool(subgraphs_stream_enabled),
-            version="v2"
+            subgraphs=bool(subgraphs_stream_enabled),
+            version="v2",
         )
 
+        stream = self.graph.astream_events(**kwargs)
+
         return {
             "stream": stream,
             "state": state,
@@ -369,12 +372,14 @@ async def prepare_regenerate_stream( # pylint: disable=too-many-arguments
 
         stream_input = self.langgraph_default_merge_state(time_travel_checkpoint.values, [message_checkpoint], input)
         subgraphs_stream_enabled = input.forwarded_props.get('stream_subgraphs') if input.forwarded_props else False
-        stream = self.graph.astream_events(
-            stream_input,
-            fork,
-            subgraps=bool(subgraphs_stream_enabled),
-            version="v2"
+
+        kwargs = self.get_stream_kwargs(
+            input=stream_input,
+            fork=fork,
+            subgraphs=bool(subgraphs_stream_enabled),
+            version="v2",
         )
+        stream = self.graph.astream_events(**kwargs)
 
         return {
             "stream": stream,
@@ -401,17 +406,25 @@ def get_schema_keys(self, config) -> SchemaKeys:
             input_schema_keys = list(input_schema["properties"].keys()) if "properties" in input_schema else []
             output_schema_keys = list(output_schema["properties"].keys()) if "properties" in output_schema else []
             config_schema_keys = list(config_schema["properties"].keys()) if "properties" in config_schema else []
+            context_schema_keys = []
+
+            if hasattr(self.graph, "context_schema") and self.graph.context_schema is not None:
+                context_schema = self.graph.context_schema().schema()
+                context_schema_keys = list(context_schema["properties"].keys()) if "properties" in context_schema else []
+
 
             return {
                 "input": [*input_schema_keys, *self.constant_schema_keys],
                 "output": [*output_schema_keys, *self.constant_schema_keys],
                 "config": config_schema_keys,
+                "context": context_schema_keys,
             }
         except Exception:
             return {
                 "input": self.constant_schema_keys,
                 "output": self.constant_schema_keys,
                 "config": [],
+                "context": [],
             }
 
     def langgraph_default_merge_state(self, state: State, messages: List[BaseMessage], input: RunAgentInput) -> State:
@@ -744,3 +757,38 @@ def end_step(self):
         self.active_run["node_name"] = None
         self.active_step = None
         return dispatch
+
+    # Check if some kwargs are enabled per LG version, to "catch all versions" and backwards compatibility
+    def get_stream_kwargs(
+            self,
+            input: Any,
+            subgraphs: bool = False,
+            version: Literal["v1", "v2"] = "v2",
+            config: Optional[RunnableConfig] = None,
+            context: Optional[Dict[str, Any]] = None,
+            fork: Optional[Any] = None,
+    ):
+        kwargs = dict(
+            input=input,
+            subgraphs=subgraphs,
+            version=version,
+        )
+
+        # Only add context if supported
+        sig = inspect.signature(self.graph.astream_events)
+        if 'context' in sig.parameters:
+            base_context = {}
+            if isinstance(config, dict) and 'configurable' in config and isinstance(config['configurable'], dict):
+                base_context.update(config['configurable'])
+            if context:  # context might be None or {}
+                base_context.update(context)
+            if base_context:  # only add if there's something to pass
+                kwargs['context'] = base_context
+
+        if config:
+            kwargs['config'] = config
+
+        if fork:
+            kwargs.update(fork)
+
+        return kwargs

typescript-sdk/integrations/langgraph/python/ag_ui_langgraph/types.py

@@ -25,7 +25,8 @@ class CustomEventNames(str, Enum):
 SchemaKeys = TypedDict("SchemaKeys", {
     "input": NotRequired[Optional[List[str]]],
     "output": NotRequired[Optional[List[str]]],
-    "config": NotRequired[Optional[List[str]]]
+    "config": NotRequired[Optional[List[str]]],
+    "context": NotRequired[Optional[List[str]]]
 })
 
 ThinkingProcess = TypedDict("ThinkingProcess", {

typescript-sdk/integrations/langgraph/python/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "ag-ui-langgraph"
-version = "0.0.11"
+version = "0.0.12-alpha.1"
 description = "Implementation of the AG-UI protocol for LangGraph."
 authors = ["Ran Shem Tov <[email protected]>"]
 readme = "README.md"