feat(langgraph): new context api (replacing config['configurable'] and config_schema) (#5243)

docs/docs/tutorials/tot/tot.ipynb

@@ -280,8 +280,8 @@
     "from typing import Optional, Dict, Any\n",
     "from typing_extensions import Annotated, TypedDict\n",
     "from langgraph.graph import StateGraph\n",
+    "from langgraph.types import Runtime\n",
     "\n",
-    "from langchain_core.runnables import RunnableConfig\n",
     "from langgraph.checkpoint.memory import MemorySaver\n",
     "from langgraph.types import Send\n",
     "\n",
@@ -307,22 +307,25 @@
     "    depth: Annotated[int, operator.add]\n",
     "\n",
     "\n",
-    "class Configuration(TypedDict, total=False):\n",
+    "class Context(TypedDict, total=False):\n",
     "    max_depth: int\n",
     "    threshold: float\n",
     "    k: int\n",
     "    beam_size: int\n",
     "\n",
+    "class EnsuredContext(TypedDict):\n",
+    "    max_depth: int\n",
+    "    threshold: float\n",
+    "    k: int\n",
+    "    beam_size: int\n",
     "\n",
-    "def _ensure_configurable(config: RunnableConfig) -> Configuration:\n",
+    "def _ensure_context(ctx: Context) -> EnsuredContext:\n",
     "    \"\"\"Get params that configure the search algorithm.\"\"\"\n",
-    "    configurable = config.get(\"configurable\", {})\n",
     "    return {\n",
-    "        **configurable,\n",
-    "        \"max_depth\": configurable.get(\"max_depth\", 10),\n",
-    "        \"threshold\": config.get(\"threshold\", 0.9),\n",
-    "        \"k\": configurable.get(\"k\", 5),\n",
-    "        \"beam_size\": configurable.get(\"beam_size\", 3),\n",
+    "        \"max_depth\": ctx.get(\"max_depth\", 10),\n",
+    "        \"threshold\": ctx.get(\"threshold\", 0.9),\n",
+    "        \"k\": ctx.get(\"k\", 5),\n",
+    "        \"beam_size\": ctx.get(\"beam_size\", 3)\n",
     "    }\n",
     "\n",
     "\n",
@@ -330,9 +333,9 @@
     "    seed: Optional[Candidate]\n",
     "\n",
     "\n",
-    "def expand(state: ExpansionState, *, config: RunnableConfig) -> Dict[str, List[str]]:\n",
+    "def expand(state: ExpansionState, *, runtime: Runtime[Context]) -> Dict[str, List[Candidate]]:\n",
     "    \"\"\"Generate the next state.\"\"\"\n",
-    "    configurable = _ensure_configurable(config)\n",
+    "    ctx = _ensure_context(runtime.context)\n",
     "    if not state.get(\"seed\"):\n",
     "        candidate_str = \"\"\n",
     "    else:\n",
@@ -342,9 +345,8 @@
     "            {\n",
     "                \"problem\": state[\"problem\"],\n",
     "                \"candidate\": candidate_str,\n",
-    "                \"k\": configurable[\"k\"],\n",
+    "                \"k\": ctx[\"k\"],\n",
     "            },\n",
-    "            config=config,\n",
     "        )\n",
     "    except Exception:\n",
     "        return {\"candidates\": []}\n",
@@ -354,7 +356,7 @@
     "    return {\"candidates\": new_candidates}\n",
     "\n",
     "\n",
-    "def score(state: ToTState) -> Dict[str, List[float]]:\n",
+    "def score(state: ToTState) -> Dict[str, Any]:\n",
     "    \"\"\"Evaluate the candidate generations.\"\"\"\n",
     "    candidates = state[\"candidates\"]\n",
     "    scored = []\n",
@@ -364,10 +366,10 @@
     "\n",
     "\n",
     "def prune(\n",
-    "    state: ToTState, *, config: RunnableConfig\n",
-    ") -> Dict[str, List[Dict[str, Any]]]:\n",
+    "    state: ToTState, *, runtime: Runtime[Context]\n",
+    ") -> Dict[str, Any]:\n",
     "    scored_candidates = state[\"scored_candidates\"]\n",
-    "    beam_size = _ensure_configurable(config)[\"beam_size\"]\n",
+    "    beam_size = _ensure_context(runtime.context)[\"beam_size\"]\n",
     "    organized = sorted(\n",
     "        scored_candidates, key=lambda candidate: candidate[1], reverse=True\n",
     "    )\n",
@@ -383,11 +385,11 @@
     "\n",
     "\n",
     "def should_terminate(\n",
-    "    state: ToTState, config: RunnableConfig\n",
+    "    state: ToTState, runtime: Runtime[Context]\n",
     ") -> Union[Literal[\"__end__\"], Send]:\n",
-    "    configurable = _ensure_configurable(config)\n",
-    "    solved = state[\"candidates\"][0].score >= configurable[\"threshold\"]\n",
-    "    if solved or state[\"depth\"] >= configurable[\"max_depth\"]:\n",
+    "    ctx = _ensure_context(runtime.context)\n",
+    "    solved = state[\"candidates\"][0].score >= ctx[\"threshold\"]\n",
+    "    if solved or state[\"depth\"] >= ctx[\"max_depth\"]:\n",
     "        return \"__end__\"\n",
     "    return [\n",
     "        Send(\"expand\", {**state, \"somevalseed\": candidate})\n",
@@ -396,7 +398,7 @@
     "\n",
     "\n",
     "# Create the graph\n",
-    "builder = StateGraph(state_schema=ToTState, config_schema=Configuration)\n",
+    "builder = StateGraph(state_schema=ToTState, context_schema=Context)\n",
     "\n",
     "# Add nodes\n",
     "builder.add_node(expand)\n",
@@ -467,13 +469,7 @@
     }
    ],
    "source": [
-    "config = {\n",
-    "    \"configurable\": {\n",
-    "        \"thread_id\": \"test_1\",\n",
-    "        \"depth\": 10,\n",
-    "    }\n",
-    "}\n",
-    "for step in graph.stream({\"problem\": puzzles[42]}, config):\n",
+    "for step in graph.stream({\"problem\": puzzles[42]}, config={\"configurable\": {\"thread_id\": \"test_1\"}}, context={\"depth\": 10}):\n",
     "    print(step)"
    ]
   },
@@ -491,7 +487,7 @@
     }
    ],
    "source": [
-    "final_state = graph.get_state(config)\n",
+    "final_state = graph.get_state({'configurable': {'thread_id': 'test_1'}})\n",
     "winning_solution = final_state.values[\"candidates\"][0]\n",
     "search_depth = final_state.values[\"depth\"]\n",
     "if winning_solution[1] == 1:\n",