feat: Migrate to LangGraph Context API

pyproject.toml

@@ -7,16 +7,16 @@ authors = [
 ]
 license = { text = "MIT" }
 readme = "README.md"
-requires-python = ">=3.9"
+requires-python = ">=3.11"
 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",
     "langchain-fireworks>=0.1.7",
     "python-dotenv>=1.0.1",
-    "langchain-elasticsearch>=0.2.2,<0.3.0",
-    "langchain-pinecone>=0.1.3,<0.2.0",
+    "langchain-elasticsearch>=0.3.0,<0.4.0",
+    "langchain-pinecone>=0.2.0,<0.3.0",
     "msgspec>=0.18.6",
     "langchain-mongodb>=0.1.9",
     "langchain-cohere>=0.2.4",

src/retrieval_graph/configuration.py

@@ -2,10 +2,9 @@
 
 from __future__ import annotations
 
-from dataclasses import dataclass, field, fields
-from typing import Annotated, Any, Literal, Optional, Type, TypeVar
-
-from langchain_core.runnables import RunnableConfig, ensure_config
+import os
+from dataclasses import dataclass, field
+from typing import Annotated, Any, Literal
 
 from retrieval_graph import prompts
 
@@ -19,7 +18,9 @@ class IndexConfiguration:
     retriever provider choice, and search parameters.
     """
 
-    user_id: str = field(metadata={"description": "Unique identifier for the user."})
+    user_id: str = field(
+        default="", metadata={"description": "Unique identifier for the user."}
+    )
 
     embedding_model: Annotated[
         str,
@@ -48,26 +49,19 @@ class IndexConfiguration:
         },
     )
 
-    @classmethod
-    def from_runnable_config(
-        cls: Type[T], config: Optional[RunnableConfig] = None
-    ) -> T:
-        """Create an IndexConfiguration instance from a RunnableConfig object.
-
-        Args:
-            cls (Type[T]): The class itself.
-            config (Optional[RunnableConfig]): The configuration object to use.
+    def __post_init__(self) -> None:
+        """Populate fields from environment variables if not already set."""
+        # Only populate from environment variables if the field is not already set
+        if not self.user_id:
+            self.user_id = os.environ.get("USER_ID", "")
 
-        Returns:
-            T: An instance of IndexConfiguration with the specified configuration.
-        """
-        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})
+        if self.embedding_model == "openai/text-embedding-3-small":
+            self.embedding_model = os.environ.get(
+                "EMBEDDING_MODEL", "openai/text-embedding-3-small"
+            )
 
-
-T = TypeVar("T", bound=IndexConfiguration)
+        if self.retriever_provider == "elastic":
+            self.retriever_provider = os.environ.get("RETRIEVER_PROVIDER", "elastic")  # type: ignore
 
 
 @dataclass(kw_only=True)
@@ -99,3 +93,19 @@ class Configuration(IndexConfiguration):
             "description": "The language model used for processing and refining queries. Should be in the form: provider/model-name."
         },
     )
+
+    def __post_init__(self) -> None:
+        """Populate fields from environment variables if not already set."""
+        # Call parent's __post_init__ first
+        super().__post_init__()
+
+        # Only populate from environment variables if the field is using the default value
+        if self.response_model == "anthropic/claude-3-5-sonnet-20240620":
+            self.response_model = os.environ.get(
+                "RESPONSE_MODEL", "anthropic/claude-3-5-sonnet-20240620"
+            )
+
+        if self.query_model == "anthropic/claude-3-haiku-20240307":
+            self.query_model = os.environ.get(
+                "QUERY_MODEL", "anthropic/claude-3-haiku-20240307"
+            )

src/retrieval_graph/graph.py

@@ -13,8 +13,8 @@
 from langchain_core.messages import BaseMessage
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_core.pydantic_v1 import BaseModel
-from langchain_core.runnables import RunnableConfig
 from langgraph.graph import StateGraph
+from langgraph.runtime import Runtime
 
 from retrieval_graph import retrieval
 from retrieval_graph.configuration import Configuration
@@ -31,7 +31,7 @@ class SearchQuery(BaseModel):
 
 
 async def generate_query(
-    state: State, *, config: RunnableConfig
+    state: State, *, runtime: Runtime[Configuration]
 ) -> dict[str, list[str]]:
     """Generate a search query based on the current state and configuration.
 
@@ -41,7 +41,7 @@ async def generate_query(
 
     Args:
         state (State): The current state containing messages and other information.
-        config (RunnableConfig | None, optional): Configuration for the query generation process.
+        runtime (Runtime[Configuration]): Runtime context containing configuration.
 
     Returns:
         dict[str, list[str]]: A dictionary with a 'queries' key containing a list of generated queries.
@@ -57,7 +57,7 @@ async def generate_query(
         human_input = get_message_text(messages[-1])
         return {"queries": [human_input]}
     else:
-        configuration = Configuration.from_runnable_config(config)
+        configuration = runtime.context
         # Feel free to customize the prompt, model, and other logic!
         prompt = ChatPromptTemplate.from_messages(
             [
@@ -74,17 +74,16 @@ async def generate_query(
                 "messages": state.messages,
                 "queries": "\n- ".join(state.queries),
                 "system_time": datetime.now(tz=timezone.utc).isoformat(),
-            },
-            config,
+            }
         )
-        generated = cast(SearchQuery, await model.ainvoke(message_value, config))
+        generated = cast(SearchQuery, await model.ainvoke(message_value))
         return {
             "queries": [generated.query],
         }
 
 
 async def retrieve(
-    state: State, *, config: RunnableConfig
+    state: State, *, runtime: Runtime[Configuration]
 ) -> dict[str, list[Document]]:
     """Retrieve documents based on the latest query in the state.
 
@@ -94,22 +93,22 @@ async def retrieve(
 
     Args:
         state (State): The current state containing queries and the retriever.
-        config (RunnableConfig | None, optional): Configuration for the retrieval process.
+        runtime (Runtime[Configuration]): Runtime context containing configuration.
 
     Returns:
         dict[str, list[Document]]: A dictionary with a single key "retrieved_docs"
         containing a list of retrieved Document objects.
     """
-    with retrieval.make_retriever(config) as retriever:
-        response = await retriever.ainvoke(state.queries[-1], config)
+    with retrieval.make_retriever(runtime) as retriever:
+        response = await retriever.ainvoke(state.queries[-1])
         return {"retrieved_docs": response}
 
 
 async def respond(
-    state: State, *, config: RunnableConfig
+    state: State, *, runtime: Runtime[Configuration]
 ) -> dict[str, list[BaseMessage]]:
     """Call the LLM powering our "agent"."""
-    configuration = Configuration.from_runnable_config(config)
+    configuration = runtime.context
     # Feel free to customize the prompt, model, and other logic!
     prompt = ChatPromptTemplate.from_messages(
         [
@@ -125,18 +124,17 @@ async def respond(
             "messages": state.messages,
             "retrieved_docs": retrieved_docs,
             "system_time": datetime.now(tz=timezone.utc).isoformat(),
-        },
-        config,
+        }
     )
-    response = await model.ainvoke(message_value, config)
+    response = await model.ainvoke(message_value)
     # We return a list, because this will get added to the existing list
     return {"messages": [response]}
 
 
 # Define a new graph (It's just a pipe)
 
 
-builder = StateGraph(State, input=InputState, config_schema=Configuration)
+builder = StateGraph(State, input=InputState, context_schema=Configuration)
 
 builder.add_node(generate_query)
 builder.add_node(retrieve)

src/retrieval_graph/index_graph.py

@@ -1,28 +1,28 @@
 """This "graph" simply exposes an endpoint for a user to upload docs to be indexed."""
 
-from typing import Optional, Sequence
+from typing import Sequence
 
 from langchain_core.documents import Document
-from langchain_core.runnables import RunnableConfig
 from langgraph.graph import StateGraph
+from langgraph.runtime import Runtime
 
 from retrieval_graph import retrieval
 from retrieval_graph.configuration import IndexConfiguration
 from retrieval_graph.state import IndexState
 
 
 def ensure_docs_have_user_id(
-    docs: Sequence[Document], config: RunnableConfig
+    docs: Sequence[Document], runtime: Runtime[IndexConfiguration]
 ) -> list[Document]:
     """Ensure that all documents have a user_id in their metadata.
 
         docs (Sequence[Document]): A sequence of Document objects to process.
-        config (RunnableConfig): A configuration object containing the user_id.
+        runtime (Runtime[IndexConfiguration]): Runtime context containing configuration.
 
     Returns:
         list[Document]: A new list of Document objects with updated metadata.
     """
-    user_id = config["configurable"]["user_id"]
+    user_id = runtime.context.user_id
     return [
         Document(
             page_content=doc.page_content, metadata={**doc.metadata, "user_id": user_id}
@@ -32,7 +32,7 @@ def ensure_docs_have_user_id(
 
 
 async def index_docs(
-    state: IndexState, *, config: Optional[RunnableConfig] = None
+    state: IndexState, *, runtime: Runtime[IndexConfiguration]
 ) -> dict[str, str]:
     """Asynchronously index documents in the given state using the configured retriever.
 
@@ -42,12 +42,10 @@ async def index_docs(
 
     Args:
         state (IndexState): The current state containing documents and retriever.
-        config (Optional[RunnableConfig]): Configuration for the indexing process.r
+        runtime (Runtime[IndexConfiguration]): Runtime context containing configuration.
     """
-    if not config:
-        raise ValueError("Configuration required to run index_docs.")
-    with retrieval.make_retriever(config) as retriever:
-        stamped_docs = ensure_docs_have_user_id(state.docs, config)
+    with retrieval.make_retriever(runtime) as retriever:
+        stamped_docs = ensure_docs_have_user_id(state.docs, runtime)
 
         await retriever.aadd_documents(stamped_docs)
     return {"docs": "delete"}
@@ -56,7 +54,7 @@ async def index_docs(
 # Define a new graph
 
 
-builder = StateGraph(IndexState, config_schema=IndexConfiguration)
+builder = StateGraph(IndexState, context_schema=IndexConfiguration)
 builder.add_node(index_docs)
 builder.add_edge("__start__", "index_docs")
 # Finally, we compile it!

src/retrieval_graph/retrieval.py

@@ -11,8 +11,8 @@
 from typing import Generator
 
 from langchain_core.embeddings import Embeddings
-from langchain_core.runnables import RunnableConfig
 from langchain_core.vectorstores import VectorStoreRetriever
+from langgraph.runtime import get_runtime
 
 from retrieval_graph.configuration import Configuration, IndexConfiguration
 
@@ -105,11 +105,9 @@ def make_mongodb_retriever(
 
 
 @contextmanager
-def make_retriever(
-    config: RunnableConfig,
-) -> Generator[VectorStoreRetriever, None, None]:
+def make_retriever() -> Generator[VectorStoreRetriever, None, None]:
     """Create a retriever for the agent, based on the current configuration."""
-    configuration = IndexConfiguration.from_runnable_config(config)
+    configuration = get_runtime(IndexConfiguration).context
     embedding_model = make_text_encoder(configuration.embedding_model)
     user_id = configuration.user_id
     if not user_id:

tests/integration_tests/test_graph.py

@@ -1,7 +1,6 @@
 import uuid
 
 import pytest
-from langchain_core.runnables import RunnableConfig
 from langsmith import expect, unit
 
 from retrieval_graph import graph, index_graph
@@ -14,27 +13,23 @@ async def test_retrieval_graph() -> None:
     user_id = "test__" + uuid.uuid4().hex
     other_user_id = "test__" + uuid.uuid4().hex
 
-    config = RunnableConfig(
-        configurable={"user_id": user_id, "retriever_provider": "elastic-local"}
-    )
+    context = {"user_id": user_id, "retriever_provider": "elastic-local"}
 
-    result = await index_graph.ainvoke({"docs": simple_doc}, config)
+    result = await index_graph.ainvoke({"docs": simple_doc}, context=context)
     expect(result["docs"]).against(lambda x: not x)  # we delete after the end
 
     res = await graph.ainvoke(
         {"messages": [("user", "Where do cats perform synchronized swimming routes?")]},
-        config,
+        context=context,
     )
     response = str(res["messages"][-1].content)
     expect(response.lower()).to_contain("bowl")
 
     res = await graph.ainvoke(
         {"messages": [("user", "Where do cats perform synchronized swimming routes?")]},
-        {
-            "configurable": {
-                "user_id": other_user_id,
-                "retriever_provider": "elastic-local",
-            }
+        context={
+            "user_id": other_user_id,
+            "retriever_provider": "elastic-local",
         },
     )
     response = str(res["messages"][-1].content)