Optional 2nd arg to StateGraph with config schema

langgraph/graph/state.py

@@ -32,10 +32,13 @@ class StateGraph(Graph):
     The signature of a reducer function is (Value, Value) -> Value.
     """
 
-    def __init__(self, schema: Type[Any]) -> None:
+    def __init__(
+        self, state_schema: Type[Any], config_schema: Optional[Type[Any]] = None
+    ) -> None:
         super().__init__()
-        self.schema = schema
-        self.channels = _get_channels(schema)
+        self.schema = state_schema
+        self.config_schema = config_schema
+        self.channels = _get_channels(state_schema)
         if any(isinstance(c, BinaryOperatorAggregate) for c in self.channels.values()):
             self.support_multiple_edges = True
         self.waiting_edges: set[tuple[tuple[str, ...], str]] = set()
@@ -134,6 +137,7 @@ class StateGraph(Graph):
 
         compiled = CompiledStateGraph(
             builder=self,
+            config_type=self.config_schema,
             nodes={},
             channels={**self.channels, START: EphemeralValue(self.schema)},
             input_channels=START,

langgraph/pregel/__init__.py

@@ -17,6 +17,7 @@ from typing import (
     Type,
     Union,
     cast,
+    get_type_hints,
     overload,
 )
 
@@ -211,6 +212,8 @@ class Pregel(
 
     checkpointer: Optional[BaseCheckpointSaver] = None
 
+    config_type: Optional[Type[Any]] = None
+
     name: str = "LangGraph"
 
     class Config:
@@ -265,6 +268,14 @@ class Pregel(
                     if self.checkpointer is not None
                     else []
                 )
+                + (
+                    [
+                        ConfigurableFieldSpec(id=name, annotation=typ)
+                        for name, typ in get_type_hints(self.config_type).items()
+                    ]
+                    if self.config_type is not None
+                    else []
+                )
             )
             # these are provided by the Pregel class
             if spec.id not in [CONFIG_KEY_READ, CONFIG_KEY_SEND]

tests/__snapshots__/test_pregel.ambr

@@ -1492,3 +1492,6 @@
   
   '''
 # ---
+# name: test_state_graph_w_config
+  '{"title": "LangGraphConfig", "type": "object", "properties": {"configurable": {"$ref": "#/definitions/Configurable"}}, "definitions": {"Configurable": {"title": "Configurable", "type": "object", "properties": {"tools": {"title": "Tools", "type": "array", "items": {"type": "string"}}}}}}'
+# ---

tests/test_pregel.py

@@ -2232,6 +2232,92 @@ def test_conditional_state_graph(
     ]
 
 
+def test_state_graph_w_config(snapshot: SnapshotAssertion) -> None:
+    from langchain.llms.fake import FakeStreamingListLLM
+    from langchain_community.tools import tool
+    from langchain_core.agents import AgentAction, AgentFinish
+    from langchain_core.prompts import PromptTemplate
+
+    class AgentState(TypedDict, total=False):
+        input: str
+        agent_outcome: Optional[Union[AgentAction, AgentFinish]]
+        intermediate_steps: Annotated[list[tuple[AgentAction, str]], operator.add]
+
+    class Config(TypedDict, total=False):
+        tools: list[str]
+
+    # Assemble the tools
+    @tool()
+    def search_api(query: str) -> str:
+        """Searches the API for the query."""
+        return f"result for {query}"
+
+    tools = [search_api]
+
+    # Construct the agent
+    prompt = PromptTemplate.from_template("Hello!")
+
+    llm = FakeStreamingListLLM(
+        responses=[
+            "tool:search_api:query",
+            "tool:search_api:another",
+            "finish:answer",
+        ]
+    )
+
+    def agent_parser(input: str) -> dict[str, Union[AgentAction, AgentFinish]]:
+        if input.startswith("finish"):
+            _, answer = input.split(":")
+            return {
+                "agent_outcome": AgentFinish(
+                    return_values={"answer": answer}, log=input
+                )
+            }
+        else:
+            _, tool_name, tool_input = input.split(":")
+            return {
+                "agent_outcome": AgentAction(
+                    tool=tool_name, tool_input=tool_input, log=input
+                )
+            }
+
+    agent = prompt | llm | agent_parser
+
+    # Define tool execution logic
+    def execute_tools(data: AgentState) -> dict:
+        agent_action: AgentAction = data.pop("agent_outcome")
+        observation = {t.name: t for t in tools}[agent_action.tool].invoke(
+            agent_action.tool_input
+        )
+        return {"intermediate_steps": [(agent_action, observation)]}
+
+    # Define decision-making logic
+    def should_continue(data: AgentState) -> str:
+        # Logic to decide whether to continue in the loop or exit
+        if isinstance(data["agent_outcome"], AgentFinish):
+            return "exit"
+        else:
+            return "continue"
+
+    # Define a new graph
+    workflow = StateGraph(AgentState, Config)
+
+    workflow.add_node("agent", agent)
+    workflow.add_node("tools", execute_tools)
+
+    workflow.set_entry_point("agent")
+
+    workflow.add_conditional_edges(
+        "agent", should_continue, {"continue": "tools", "exit": END}
+    )
+
+    workflow.add_edge("tools", "agent")
+
+    app = workflow.compile()
+
+    assert app.config_schema().schema_json() == snapshot
+
+
 def test_conditional_entrypoint_graph_state(snapshot: SnapshotAssertion) -> None:
     class AgentState(TypedDict, total=False):
         input: str