chore: use new context API (#23)

Author: sydney-runkle

README.md

@@ -36,10 +36,6 @@ cp .env.example .env
 
 The primary [search tool](./src/react_agent/tools.py) [^1] used is [Tavily](https://tavily.com/). Create an API key [here](https://app.tavily.com/sign-in).
 
-<!--
-Setup instruction auto-generated by `langgraph template lock`. DO NOT EDIT MANUALLY.
--->
-
 ### Setup Model
 
 The defaults values for `model` are shown below:
@@ -70,23 +66,14 @@ To use OpenAI's chat models:
 OPENAI_API_KEY=your-api-key
 ```
 
-
-
-
-
-<!--
-End setup instructions
--->
-
-
 3. Customize whatever you'd like in the code.
 4. Open the folder LangGraph Studio!
 
 ## How to customize
 
 1. **Add new tools**: Extend the agent's capabilities by adding new tools in [tools.py](./src/react_agent/tools.py). These can be any Python functions that perform specific tasks.
-2. **Select a different model**: We default to Anthropic's Claude 3 Sonnet. You can select a compatible chat model using `provider/model-name` via configuration. Example: `openai/gpt-4-turbo-preview`.
-3. **Customize the prompt**: We provide a default system prompt in [prompts.py](./src/react_agent/prompts.py). You can easily update this via configuration in the studio.
+2. **Select a different model**: We default to Anthropic's Claude 3 Sonnet. You can select a compatible chat model using `provider/model-name` via runtime context. Example: `openai/gpt-4-turbo-preview`.
+3. **Customize the prompt**: We provide a default system prompt in [prompts.py](./src/react_agent/prompts.py). You can easily update this via context in the studio.
 
 You can also quickly extend this template by:
 
@@ -95,7 +82,7 @@ You can also quickly extend this template by:
 
 ## Development
 
-While iterating on your graph, you can edit past state and rerun your app from past states to debug specific nodes. Local changes will be automatically applied via hot reload. Try adding an interrupt before the agent calls tools, updating the default system message in `src/react_agent/configuration.py` to take on a persona, or adding additional nodes and edges!
+While iterating on your graph, you can edit past state and rerun your app from past states to debug specific nodes. Local changes will be automatically applied via hot reload. Try adding an interrupt before the agent calls tools, updating the default system message in `src/react_agent/context.py` to take on a persona, or adding additional nodes and edges!
 
 Follow up requests will be appended to the same thread. You can create an entirely new thread, clearing previous history, using the `+` button in the top right.
 
@@ -104,138 +91,3 @@ You can find the latest (under construction) docs on [LangGraph](https://github.
 LangGraph Studio also integrates with [LangSmith](https://smith.langchain.com/) for more in-depth tracing and collaboration with teammates.
 
 [^1]: https://python.langchain.com/docs/concepts/#tools
-
-<!--
-Configuration auto-generated by `langgraph template lock`. DO NOT EDIT MANUALLY.
-{
-  "config_schemas": {
-    "agent": {
-      "type": "object",
-      "properties": {
-        "model": {
-          "type": "string",
-          "default": "anthropic/claude-3-5-sonnet-20240620",
-          "description": "The name of the language model to use for the agent's main interactions. Should be in the form: provider/model-name.",
-          "environment": [
-            {
-              "value": "anthropic/claude-1.2",
-              "variables": "ANTHROPIC_API_KEY"
-            },
-            {
-              "value": "anthropic/claude-2.0",
-              "variables": "ANTHROPIC_API_KEY"
-            },
-            {
-              "value": "anthropic/claude-2.1",
-              "variables": "ANTHROPIC_API_KEY"
-            },
-            {
-              "value": "anthropic/claude-3-5-sonnet-20240620",
-              "variables": "ANTHROPIC_API_KEY"
-            },
-            {
-              "value": "anthropic/claude-3-haiku-20240307",
-              "variables": "ANTHROPIC_API_KEY"
-            },
-            {
-              "value": "anthropic/claude-3-opus-20240229",
-              "variables": "ANTHROPIC_API_KEY"
-            },
-            {
-              "value": "anthropic/claude-3-sonnet-20240229",
-              "variables": "ANTHROPIC_API_KEY"
-            },
-            {
-              "value": "anthropic/claude-instant-1.2",
-              "variables": "ANTHROPIC_API_KEY"
-            },
-            {
-              "value": "openai/gpt-3.5-turbo",
-              "variables": "OPENAI_API_KEY"
-            },
-            {
-              "value": "openai/gpt-3.5-turbo-0125",
-              "variables": "OPENAI_API_KEY"
-            },
-            {
-              "value": "openai/gpt-3.5-turbo-0301",
-              "variables": "OPENAI_API_KEY"
-            },
-            {
-              "value": "openai/gpt-3.5-turbo-0613",
-              "variables": "OPENAI_API_KEY"
-            },
-            {
-              "value": "openai/gpt-3.5-turbo-1106",
-              "variables": "OPENAI_API_KEY"
-            },
-            {
-              "value": "openai/gpt-3.5-turbo-16k",
-              "variables": "OPENAI_API_KEY"
-            },
-            {
-              "value": "openai/gpt-3.5-turbo-16k-0613",
-              "variables": "OPENAI_API_KEY"
-            },
-            {
-              "value": "openai/gpt-4",
-              "variables": "OPENAI_API_KEY"
-            },
-            {
-              "value": "openai/gpt-4-0125-preview",
-              "variables": "OPENAI_API_KEY"
-            },
-            {
-              "value": "openai/gpt-4-0314",
-              "variables": "OPENAI_API_KEY"
-            },
-            {
-              "value": "openai/gpt-4-0613",
-              "variables": "OPENAI_API_KEY"
-            },
-            {
-              "value": "openai/gpt-4-1106-preview",
-              "variables": "OPENAI_API_KEY"
-            },
-            {
-              "value": "openai/gpt-4-32k",
-              "variables": "OPENAI_API_KEY"
-            },
-            {
-              "value": "openai/gpt-4-32k-0314",
-              "variables": "OPENAI_API_KEY"
-            },
-            {
-              "value": "openai/gpt-4-32k-0613",
-              "variables": "OPENAI_API_KEY"
-            },
-            {
-              "value": "openai/gpt-4-turbo",
-              "variables": "OPENAI_API_KEY"
-            },
-            {
-              "value": "openai/gpt-4-turbo-preview",
-              "variables": "OPENAI_API_KEY"
-            },
-            {
-              "value": "openai/gpt-4-vision-preview",
-              "variables": "OPENAI_API_KEY"
-            },
-            {
-              "value": "openai/gpt-4o",
-              "variables": "OPENAI_API_KEY"
-            },
-            {
-              "value": "openai/gpt-4o-mini",
-              "variables": "OPENAI_API_KEY"
-            }
-          ]
-        }
-      },
-      "environment": [
-        "TAVILY_API_KEY"
-      ]
-    }
-  }
-}
--->

pyproject.toml

@@ -9,7 +9,7 @@ readme = "README.md"
 license = { text = "MIT" }
 requires-python = ">=3.11,<4.0"
 dependencies = [
-    "langgraph>=0.2.6",
+    "langgraph>=0.6.0,<0.7.0",
     "langchain-openai>=0.1.22",
     "langchain-anthropic>=0.1.23",
     "langchain>=0.2.14",

src/react_agent/configuration.py renamed to src/react_agent/context.py

@@ -2,18 +2,16 @@
 
 from __future__ import annotations
 
+import os
 from dataclasses import dataclass, field, fields
 from typing import Annotated
 
-from langchain_core.runnables import ensure_config
-from langgraph.config import get_config
-
-from react_agent import prompts
+from . import prompts
 
 
 @dataclass(kw_only=True)
-class Configuration:
-    """The configuration for the agent."""
+class Context:
+    """The context for the agent."""
 
     system_prompt: str = field(
         default=prompts.SYSTEM_PROMPT,
@@ -38,14 +36,11 @@ class Configuration:
         },
     )
 
-    @classmethod
-    def from_context(cls) -> Configuration:
-        """Create a Configuration instance from a RunnableConfig object."""
-        try:
-            config = get_config()
-        except RuntimeError:
-            config = None
-        config = ensure_config(config)
-        configurable = config.get("configurable") or {}
-        _fields = {f.name for f in fields(cls) if f.init}
-        return cls(**{k: v for k, v in configurable.items() if k in _fields})
+    def __post_init__(self) -> None:
+        """Fetch env vars for attributes that were not passed as args."""
+        for f in fields(self):
+            if not f.init:
+                continue
+
+            if getattr(self, f.name) == f.default:
+                setattr(self, f.name, os.environ.get(f.name.upper(), f.default))

src/react_agent/graph.py

@@ -9,16 +9,19 @@
 from langchain_core.messages import AIMessage
 from langgraph.graph import StateGraph
 from langgraph.prebuilt import ToolNode
+from langgraph.runtime import Runtime
 
-from react_agent.configuration import Configuration
+from react_agent.context import Context
 from react_agent.state import InputState, State
 from react_agent.tools import TOOLS
 from react_agent.utils import load_chat_model
 
 # Define the function that calls the model
 
 
-async def call_model(state: State) -> Dict[str, List[AIMessage]]:
+async def call_model(
+    state: State, runtime: Runtime[Context]
+) -> Dict[str, List[AIMessage]]:
     """Call the LLM powering our "agent".
 
     This function prepares the prompt, initializes the model, and processes the response.
@@ -30,13 +33,11 @@ async def call_model(state: State) -> Dict[str, List[AIMessage]]:
     Returns:
         dict: A dictionary containing the model's response message.
     """
-    configuration = Configuration.from_context()
-
     # Initialize the model with tool binding. Change the model or add more tools here.
-    model = load_chat_model(configuration.model).bind_tools(TOOLS)
+    model = load_chat_model(runtime.context.model).bind_tools(TOOLS)
 
     # Format the system prompt. Customize this to change the agent's behavior.
-    system_message = configuration.system_prompt.format(
+    system_message = runtime.context.system_prompt.format(
         system_time=datetime.now(tz=UTC).isoformat()
     )
 
@@ -65,7 +66,7 @@ async def call_model(state: State) -> Dict[str, List[AIMessage]]:
 
 # Define a new graph
 
-builder = StateGraph(State, input=InputState, config_schema=Configuration)
+builder = StateGraph(State, input_schema=InputState, context_schema=Context)
 
 # Define the two nodes we will cycle between
 builder.add_node(call_model)

src/react_agent/tools.py

@@ -8,9 +8,10 @@
 
 from typing import Any, Callable, List, Optional, cast
 
-from langchain_tavily import TavilySearch  # type: ignore[import-not-found]
+from langchain_tavily import TavilySearch
+from langgraph.runtime import get_runtime
 
-from react_agent.configuration import Configuration
+from react_agent.context import Context
 
 
 async def search(query: str) -> Optional[dict[str, Any]]:
@@ -20,8 +21,8 @@ async def search(query: str) -> Optional[dict[str, Any]]:
     to provide comprehensive, accurate, and trusted results. It's particularly useful
     for answering questions about current events.
     """
-    configuration = Configuration.from_context()
-    wrapped = TavilySearch(max_results=configuration.max_search_results)
+    runtime = get_runtime(Context)
+    wrapped = TavilySearch(max_results=runtime.context.max_search_results)
     return cast(dict[str, Any], await wrapped.ainvoke({"query": query}))
 
 

tests/integration_tests/test_graph.py

@@ -2,14 +2,15 @@
 from langsmith import unit
 
 from react_agent import graph
+from react_agent.context import Context
 
 
 @pytest.mark.asyncio
 @unit
 async def test_react_agent_simple_passthrough() -> None:
     res = await graph.ainvoke(
-        {"messages": [("user", "Who is the founder of LangChain?")]},
-        {"configurable": {"system_prompt": "You are a helpful AI assistant."}},
+        {"messages": [("user", "Who is the founder of LangChain?")]},  # type: ignore
+        context=Context(system_prompt="You are a helpful AI assistant."),
     )
 
     assert "harrison" in str(res["messages"][-1].content).lower()

tests/unit_tests/test_configuration.py

@@ -1,5 +1,20 @@
-from react_agent.configuration import Configuration
+import os
 
+from react_agent.context import Context
 
-def test_configuration_empty() -> None:
-    Configuration.from_context()
+
+def test_context_init() -> None:
+    context = Context(model="openai/gpt-4o-mini")
+    assert context.model == "openai/gpt-4o-mini"
+
+
+def test_context_init_with_env_vars() -> None:
+    os.environ["MODEL"] = "openai/gpt-4o-mini"
+    context = Context()
+    assert context.model == "openai/gpt-4o-mini"
+
+
+def test_context_init_with_env_vars_and_passed_values() -> None:
+    os.environ["MODEL"] = "openai/gpt-4o-mini"
+    context = Context(model="openai/gpt-5o-mini")
+    assert context.model == "openai/gpt-5o-mini"