mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-16 22:57:52 +02:00
118 lines
4.2 KiB
Python
118 lines
4.2 KiB
Python
#!/usr/bin/env python3
|
|
"""Test script to verify reserved keyword injection works correctly."""
|
|
|
|
from typing import Any
|
|
from langchain_core.tools import tool
|
|
from langgraph.prebuilt import ToolNode
|
|
from langgraph.runtime import Runtime
|
|
from langgraph.store.memory import InMemoryStore
|
|
from langchain_core.messages import AIMessage, ToolCall
|
|
|
|
|
|
# Test tool with reserved keyword 'state'
|
|
@tool
|
|
def tool_with_state(x: int, state) -> str:
|
|
"""Tool that uses reserved keyword 'state'."""
|
|
return f"x={x}, state_keys={list(state.keys()) if isinstance(state, dict) else 'not_dict'}"
|
|
|
|
|
|
# Test tool with reserved keyword 'runtime'
|
|
@tool
|
|
def tool_with_runtime(x: int, runtime) -> str:
|
|
"""Tool that uses reserved keyword 'runtime'."""
|
|
has_store = runtime.store is not None if hasattr(runtime, 'store') else False
|
|
return f"x={x}, has_store={has_store}"
|
|
|
|
|
|
# Test tool with both reserved keywords
|
|
@tool
|
|
def tool_with_both(x: int, state, runtime) -> str:
|
|
"""Tool that uses both reserved keywords."""
|
|
has_store = runtime.store is not None if hasattr(runtime, 'store') else False
|
|
return f"x={x}, state_keys={list(state.keys()) if isinstance(state, dict) else 'not_dict'}, has_store={has_store}"
|
|
|
|
|
|
# Test regular tool without injection
|
|
@tool
|
|
def regular_tool(x: int, y: str) -> str:
|
|
"""Regular tool without injection."""
|
|
return f"x={x}, y={y}"
|
|
|
|
|
|
def test_reserved_keywords():
|
|
"""Test that reserved keywords work correctly."""
|
|
|
|
# Create ToolNode with all test tools
|
|
tools = [tool_with_state, tool_with_runtime, tool_with_both, regular_tool]
|
|
node = ToolNode(tools)
|
|
|
|
# Check that reserved keywords are detected
|
|
print("Tool to state args:", node.tool_to_state_args)
|
|
print("Tool to runtime args:", node.tool_to_runtime_arg)
|
|
|
|
# Check tool schemas - reserved keywords should be excluded
|
|
for tool_name, tool_obj in node.tools_by_name.items():
|
|
schema = tool_obj.get_input_schema()
|
|
print(f"\n{tool_name} schema fields:", list(schema.__fields__.keys()))
|
|
|
|
# Verify reserved keywords are not in the schema
|
|
if tool_name == "tool_with_state":
|
|
assert "state" not in schema.__fields__, f"'state' should be excluded from {tool_name} schema"
|
|
elif tool_name == "tool_with_runtime":
|
|
assert "runtime" not in schema.__fields__, f"'runtime' should be excluded from {tool_name} schema"
|
|
elif tool_name == "tool_with_both":
|
|
assert "state" not in schema.__fields__, f"'state' should be excluded from {tool_name} schema"
|
|
assert "runtime" not in schema.__fields__, f"'runtime' should be excluded from {tool_name} schema"
|
|
|
|
print("\nAll schema checks passed!")
|
|
|
|
# Test actual injection
|
|
store = InMemoryStore()
|
|
state = {"messages": [], "foo": "bar"}
|
|
|
|
# Create tool calls
|
|
tool_call1: ToolCall = {
|
|
"name": "tool_with_state",
|
|
"args": {"x": 1},
|
|
"id": "1",
|
|
"type": "tool_call"
|
|
}
|
|
|
|
tool_call2: ToolCall = {
|
|
"name": "tool_with_runtime",
|
|
"args": {"x": 2},
|
|
"id": "2",
|
|
"type": "tool_call"
|
|
}
|
|
|
|
tool_call3: ToolCall = {
|
|
"name": "regular_tool",
|
|
"args": {"x": 3, "y": "test"},
|
|
"id": "3",
|
|
"type": "tool_call"
|
|
}
|
|
|
|
# Test injection
|
|
from langchain_core.runnables import RunnableConfig
|
|
config = RunnableConfig(configurable={"context": {"user_id": "test_user"}})
|
|
|
|
injected1 = node.inject_tool_args(tool_call1, state, store, config)
|
|
print(f"\nInjected args for tool_with_state: {injected1['args']}")
|
|
assert "state" in injected1["args"], "State should be injected"
|
|
|
|
injected2 = node.inject_tool_args(tool_call2, state, store, config)
|
|
print(f"Injected args for tool_with_runtime: {injected2['args']}")
|
|
assert "runtime" in injected2["args"], "Runtime should be injected"
|
|
|
|
injected3 = node.inject_tool_args(tool_call3, state, store, config)
|
|
print(f"Injected args for regular_tool: {injected3['args']}")
|
|
assert "state" not in injected3["args"], "State should not be injected"
|
|
assert "runtime" not in injected3["args"], "Runtime should not be injected"
|
|
|
|
print("\nAll injection tests passed!")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_reserved_keywords()
|
|
|