mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-31 20:29:46 +02:00
fix(prebuilt): handle injected NotRequired keys (#7392)
Resolves https://github.com/langchain-ai/langchain/issues/35585 This would previously raise KeyError: ```python from typing import Annotated from langchain_core.tools import tool from langchain.agents import create_agent from typing_extensions import NotRequired from langgraph.prebuilt import InjectedState from langchain.agents import AgentState class CustomAgentState(AgentState): city: NotRequired[str] @tool def get_weather(city: Annotated[str | None, InjectedState("city")] = None) -> str: """Get weather for a given city.""" if city is None: city = "Boston" return f"It's always sunny in {city}!" agent = create_agent( model="claude-sonnet-4-6", tools=[get_weather], system_prompt="You are a helpful assistant", state_schema=CustomAgentState, ) input_message = { "role": "user", "content": "What's the weather?", } result = agent.invoke({"messages": [input_message]}) for m in result["messages"]: m.pretty_print() ``` --------- Co-authored-by: Sydney Runkle <sydneymarierunkle@gmail.com>
This commit is contained in:
committed by
Sydney Runkle
co-authored by
Sydney Runkle
parent
792f779c18
commit
024cf29274
@@ -614,6 +614,7 @@ class _InjectedArgs:
|
||||
store: str | None
|
||||
runtime: str | None
|
||||
all_injected_keys: set[str]
|
||||
_optional_state_args: set[str]
|
||||
|
||||
|
||||
class ToolNode(RunnableCallable):
|
||||
@@ -1333,7 +1334,7 @@ class ToolNode(RunnableCallable):
|
||||
return tool_call
|
||||
|
||||
tool_call_copy: ToolCall = copy(tool_call)
|
||||
injected_args = {}
|
||||
injected_args: dict[str, Any] = {}
|
||||
|
||||
# Inject state
|
||||
if injected.state:
|
||||
@@ -1361,14 +1362,20 @@ class ToolNode(RunnableCallable):
|
||||
# Extract state values
|
||||
if isinstance(state, dict):
|
||||
for tool_arg, state_field in injected.state.items():
|
||||
injected_args[tool_arg] = (
|
||||
state[state_field] if state_field else state
|
||||
)
|
||||
if not state_field:
|
||||
injected_args[tool_arg] = state
|
||||
elif state_field in state:
|
||||
injected_args[tool_arg] = state[state_field]
|
||||
elif tool_arg not in injected._optional_state_args:
|
||||
raise KeyError(state_field)
|
||||
else:
|
||||
for tool_arg, state_field in injected.state.items():
|
||||
injected_args[tool_arg] = (
|
||||
getattr(state, state_field) if state_field else state
|
||||
)
|
||||
if not state_field:
|
||||
injected_args[tool_arg] = state
|
||||
elif hasattr(state, state_field):
|
||||
injected_args[tool_arg] = getattr(state, state_field)
|
||||
elif tool_arg not in injected._optional_state_args:
|
||||
raise AttributeError(state_field)
|
||||
|
||||
# Inject store
|
||||
if injected.store:
|
||||
@@ -1859,6 +1866,7 @@ def _get_all_injected_args(tool: BaseTool) -> _InjectedArgs:
|
||||
store_arg: str | None = None
|
||||
runtime_arg: str | None = None
|
||||
all_injected_keys: set[str] = set()
|
||||
_optional_state_args: set[str] = set()
|
||||
|
||||
for name, type_ in all_annotations.items():
|
||||
# Track all InjectedToolArg-annotated params (including custom subclasses)
|
||||
@@ -1873,6 +1881,9 @@ def _get_all_injected_args(tool: BaseTool) -> _InjectedArgs:
|
||||
if state_inj := _get_injection_from_type(type_, InjectedState):
|
||||
if isinstance(state_inj, InjectedState) and state_inj.field:
|
||||
state_args[name] = state_inj.field
|
||||
field_info = full_schema.model_fields.get(name)
|
||||
if field_info and not field_info.is_required():
|
||||
_optional_state_args.add(name)
|
||||
else:
|
||||
state_args[name] = None
|
||||
|
||||
@@ -1889,4 +1900,5 @@ def _get_all_injected_args(tool: BaseTool) -> _InjectedArgs:
|
||||
store=store_arg,
|
||||
runtime=runtime_arg,
|
||||
all_injected_keys=all_injected_keys,
|
||||
_optional_state_args=_optional_state_args,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user