From 467f42d799108cb09fa0dc9aa00e494eef2546f8 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 3 May 2024 10:05:14 -0700 Subject: [PATCH] Optional 2nd arg to StateGraph with config schema --- langgraph/graph/state.py | 10 +++- langgraph/pregel/__init__.py | 11 ++++ tests/__snapshots__/test_pregel.ambr | 3 + tests/test_pregel.py | 86 ++++++++++++++++++++++++++++ 4 files changed, 107 insertions(+), 3 deletions(-) diff --git a/langgraph/graph/state.py b/langgraph/graph/state.py index 5e7bfe8b1..85b38d585 100644 --- a/langgraph/graph/state.py +++ b/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, diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index e02de7172..85f8364c5 100644 --- a/langgraph/pregel/__init__.py +++ b/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] diff --git a/tests/__snapshots__/test_pregel.ambr b/tests/__snapshots__/test_pregel.ambr index b5bc325f6..1b1866624 100644 --- a/tests/__snapshots__/test_pregel.ambr +++ b/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"}}}}}}' +# --- diff --git a/tests/test_pregel.py b/tests/test_pregel.py index f477d4988..76358f28c 100644 --- a/tests/test_pregel.py +++ b/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