WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-17 21:25:46 +02:00
Code Issues Packages Projects Releases Wiki Activity

Tool Validator Node (#468)

Browse Source
This commit is contained in:
William FH
2024-05-15 20:02:38 -07:00
committed by GitHub
parent a308eb5f3b
commit 4209508556
7 changed files with 332 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Makefile
+1 -1
View File
@@ -55,7 +55,7 @@ build-docs:
serve-clean-docs: clean-docs
poetry run python docs/_scripts/copy_notebooks.py
poetry run python -m mkdocs serve -c -f mkdocs.yml --strict -w ./langgraph
poetry run python -m mkdocs serve -c -f docs/mkdocs.yml --strict -w ./langgraph
serve-docs:
poetry run python docs/_scripts/copy_notebooks.py
docs/docs/concepts/index.md
+1 -1
View File
@@ -51,7 +51,7 @@ We will go through a full execution of a StateGraph later, but first, lets explo
## Nodes
In StateGraph, nodes are typically python functions (sync or `async`) where the **first** positional argument is the [state](#state), and (optionally), the **second** positional argument is a "config", containing optional [configurable parameters](#configuration) (such as a `thread_id`).
In StateGraph, nodes are typically python functions (sync or `async`) where the **first** positional argument is the [state](#state-management), and (optionally), the **second** positional argument is a "config", containing optional [configurable parameters](#configuration) (such as a `thread_id`).
Similar to `NetworkX`, you add these nodes to a graph using the [add_node](https://langchain-ai.github.io/langgraph/reference/graphs/#langgraph.graph.MessageGraph) method:
docs/docs/reference/graphs.md
+1 -1
View File
@@ -1,6 +1,6 @@
# Graph Definitions
Graphs are the core abstraction of LangGraph. Each [StateGraph](#langgraph.graph.StateGraph) implementation is used to create graph workflows. Once compiled, you can run the [CompiledGraph](#compiledgraph) to run the application.
Graphs are the core abstraction of LangGraph. Each [StateGraph](#stategraph) implementation is used to create graph workflows. Once compiled, you can run the [CompiledGraph](#compiledgraph) to run the application.
## StateGraph
docs/docs/reference/prebuilt.md
+18 -8
View File
@@ -1,5 +1,13 @@
# Prebuilt
## create_react_agent
```python
from langgraph.prebuilt import create_react_agent
```
::: langgraph.prebuilt.create_react_agent
## ToolNode
```python
@@ -31,13 +39,6 @@ from langgraph.prebuilt import ToolInvocation
heading_level: 4
## create_react_agent
```python
from langgraph.prebuilt import create_react_agent
```
::: langgraph.prebuilt.create_react_agent
## `tools_condition`
@@ -45,4 +46,13 @@ from langgraph.prebuilt import create_react_agent
from langgraph.prebuilt import tools_condition
```
::: langgraph.prebuilt.tools_condition
::: langgraph.prebuilt.tools_condition
## ValidationNode
```python
from langgraph.prebuilt import ValidationNode
```
::: langgraph.prebuilt.ValidationNode
langgraph/prebuilt/__init__.py
+2
View File
@@ -4,6 +4,7 @@ from langgraph.prebuilt.agent_executor import create_agent_executor
from langgraph.prebuilt.chat_agent_executor import create_react_agent
from langgraph.prebuilt.tool_executor import ToolExecutor, ToolInvocation
from langgraph.prebuilt.tool_node import ToolNode, tools_condition
from langgraph.prebuilt.tool_validator import ValidationNode
__all__ = [
"create_agent_executor",
@@ -13,4 +14,5 @@ __all__ = [
"ToolInvocation",
"ToolNode",
"tools_condition",
"ValidationNode",
]
langgraph/prebuilt/tool_validator.py
+235
View File
@@ -0,0 +1,235 @@
"""This module provides a ValidationNode class that can be used to validate tool calls
in a langchain graph. It applies a pydantic schema to tool_calls in the models' outputs,
and returns a ToolMessage with the validated content. If the schema is not valid, it
returns a ToolMessage with the error message. The ValidationNode can be used in a
StateGraph with a "messages" key or in a MessageGraph. If multiple tool calls are
requested, they will be run in parallel.
"""
from typing import (
Any,
Callable,
Dict,
Optional,
Sequence,
Tuple,
Type,
Union,
cast,
)
from langchain_core.messages import (
AIMessage,
AnyMessage,
ToolCall,
ToolMessage,
)
from langchain_core.pydantic_v1 import BaseModel, ValidationError
from langchain_core.runnables import (
RunnableConfig,
)
from langchain_core.runnables.config import get_executor_for_config
from langchain_core.tools import BaseTool, create_schema_from_function
from pydantic import BaseModel as BaseModelV2
from pydantic import ValidationError as ValidationErrorV2
from langgraph.utils import RunnableCallable
def _default_format_error(
error: BaseException, call: ToolCall, schema: Type[BaseModel]
) -> str:
"""Default error formatting function."""
return f"{repr(error)}\n\nRespond after fixing all validation errors."
class ValidationNode(RunnableCallable):
"""A node that validates all tools requests from the last AIMessage.
It can be used either in StateGraph with a "messages" key or in MessageGraph.
!!! note
This node does not actually **run** the tools, it only validates the tool calls,
which is useful for extraction and other use cases where you need to generate
structured output that conforms to a complex schema without losing the original
messages and tool IDs (for use in multi-turn conversations).
Args:
schemas: A list of schemas to validate the tool calls with. These can be
any of the following:
- A pydantic BaseModel class
- A BaseTool instance (the args_schema will be used)
- A function (a schema will be created from the function signature)
format_error: A function that takes an exception, a ToolCall, and a schema
and returns a formatted error string. By default, it returns the
exception repr and a message to respond after fixing validation errors.
name: The name of the node.
tags: A list of tags to add to the node.
Returns:
(Union[Dict[str, List[ToolMessage]], Sequence[ToolMessage]]): A list of ToolMessages with the validated content or error messages.
Examples:
Example usage for re-prompting the model to generate a valid response:
>>> from typing import Literal
...
>>> from langchain_anthropic import ChatAnthropic
>>> from langchain_core.pydantic_v1 import BaseModel, validator
...
>>> from langgraph.graph import END, START, MessageGraph
>>> from langgraph.prebuilt import ValidationNode
...
...
>>> class SelectNumber(BaseModel):
... a: int
...
... @validator("a")
... def a_must_be_meaningful(cls, v):
... if v != 37:
... raise ValueError("Only 37 is allowed")
... return v
...
...
>>> builder = MessageGraph()
>>> llm = ChatAnthropic(model="claude-3-haiku-20240307").bind_tools([SelectNumber])
>>> builder.add_node("model", llm)
>>> builder.add_node("validation", ValidationNode([SelectNumber]))
>>> builder.add_edge(START, "model")
...
...
>>> def should_validate(state: list) -> Literal["validation", "__end__"]:
... if state[-1].tool_calls:
... return "validation"
... return END
...
...
>>> builder.add_conditional_edges("model", should_validate)
...
...
>>> def should_reprompt(state: list) -> Literal["model", "__end__"]:
... for msg in state[::-1]:
... # None of the tool calls were errors
... if msg.type == "ai":
... return END
... if msg.additional_kwargs.get("is_error"):
... return "model"
... return END
...
...
>>> builder.add_conditional_edges("validation", should_reprompt)
...
...
>>> graph = builder.compile()
>>> res = graph.invoke(("user", "Select a number, any number"))
>>> # Show the retry logic
>>> for msg in res:
... msg.pretty_print()
================================ Human Message =================================
Select a number, any number
================================== Ai Message ==================================
\n
[{'id': 'toolu_01JSjT9Pq8hGmTgmMPc6KnvM', 'input': {'a': 42}, 'name': 'SelectNumber', 'type': 'tool_use'}]
Tool Calls:
SelectNumber (toolu_01JSjT9Pq8hGmTgmMPc6KnvM)
Call ID: toolu_01JSjT9Pq8hGmTgmMPc6KnvM
Args:
a: 42
================================= Tool Message =================================
Name: SelectNumber
\n
ValidationError(model='SelectNumber', errors=[{'loc': ('a',), 'msg': 'Only 37 is allowed', 'type': 'value_error'}])
\n
Respond after fixing all validation errors.
================================== Ai Message ==================================
\n
[{'id': 'toolu_01PkxSVxNxc5wqwCPW1FiSmV', 'input': {'a': 37}, 'name': 'SelectNumber', 'type': 'tool_use'}]
Tool Calls:
SelectNumber (toolu_01PkxSVxNxc5wqwCPW1FiSmV)
Call ID: toolu_01PkxSVxNxc5wqwCPW1FiSmV
Args:
a: 37
================================= Tool Message =================================
Name: SelectNumber
\n
{"a": 37}
"""
def __init__(
self,
schemas: Sequence[Union[BaseTool, Type[BaseModel], Callable]],
*,
format_error: Optional[
Callable[[BaseException, ToolCall, Type[BaseModel]], str]
] = None,
name: str = "validation",
tags: Optional[list[str]] = None,
) -> None:
super().__init__(self._func, None, name=name, tags=tags, trace=False)
self._format_error = format_error or _default_format_error
self.schemas_by_name: Dict[str, Type[BaseModel]] = {}
for schema in schemas:
if isinstance(schema, BaseTool):
if schema.args_schema is None:
raise ValueError(
f"Tool {schema.name} does not have an args_schema defined."
)
self.schemas_by_name[schema.name] = schema.args_schema
elif isinstance(schema, type) and issubclass(
schema, (BaseModel, BaseModelV2)
):
self.schemas_by_name[schema.__name__] = cast(Type[BaseModel], schema)
elif callable(schema):
base_model = create_schema_from_function("Validation", schema)
self.schemas_by_name[schema.__name__] = base_model
else:
raise ValueError(
f"Unsupported input to ValidationNode. Expected BaseModel, tool or function. Got: {type(schema)}."
)
def _get_message(
self, input: Union[list[AnyMessage], dict[str, Any]]
) -> Tuple[str, AIMessage]:
"""Extract the last AIMessage from the input."""
if isinstance(input, list):
output_type = "list"
messages: list = input
elif messages := input.get("messages", []):
output_type = "dict"
else:
raise ValueError("No message found in input")
message: AnyMessage = messages[-1]
if not isinstance(message, AIMessage):
raise ValueError("Last message is not an AIMessage")
return output_type, message
def _func(
self, input: Union[list[AnyMessage], dict[str, Any]], config: RunnableConfig
) -> Any:
"""Validate and run tool calls synchronously."""
output_type, message = self._get_message(input)
def run_one(call: ToolCall):
schema = self.schemas_by_name[call["name"]]
try:
output = schema.validate(call["args"])
return ToolMessage(
content=output.json(),
name=call["name"],
tool_call_id=cast(str, call["id"]),
)
except (ValidationError, ValidationErrorV2) as e:
return ToolMessage(
content=self._format_error(e, call, schema),
name=call["name"],
tool_call_id=cast(str, call["id"]),
additional_kwargs={"is_error": True},
)
with get_executor_for_config(config) as executor:
outputs = [*executor.map(run_one, message.tool_calls)]
if output_type == "list":
return outputs
else:
return {"messages": outputs}
tests/test_prebuilt.py
+74 -1
View File
@@ -1,5 +1,6 @@
from typing import Any, Callable, Dict, List, Optional, Sequence, Type, Union
import pytest
from langchain_core.callbacks import (
CallbackManagerForLLMRun,
)
@@ -18,8 +19,10 @@ from langchain_core.outputs import ChatGeneration, ChatResult
from langchain_core.pydantic_v1 import BaseModel
from langchain_core.runnables import Runnable, RunnableLambda
from langchain_core.tools import BaseTool
from langchain_core.tools import tool as dec_tool
from pydantic import BaseModel as BaseModelV2
from langgraph.prebuilt import ToolNode, create_react_agent
from langgraph.prebuilt import ToolNode, ValidationNode, create_react_agent
class FakeToolCallingModel(BaseChatModel):
@@ -152,3 +155,73 @@ async def test_tool_node():
tool_message: ToolMessage = result2["messages"][-1]
assert tool_message.type == "tool"
assert tool_message.content == "tool2: 2 - bar"
def my_function(some_val: int, some_other_val: str) -> str:
return f"{some_val} - {some_other_val}"
class MyModel(BaseModel):
some_val: int
some_other_val: str
class MyModelV2(BaseModelV2):
some_val: int
some_other_val: str
@dec_tool
def my_tool(some_val: int, some_other_val: str) -> str:
"""Cool."""
return f"{some_val} - {some_other_val}"
@pytest.mark.parametrize(
"tool_schema",
[
my_function,
MyModel,
MyModelV2,
my_tool,
],
)
@pytest.mark.parametrize("use_message_key", [True, False])
async def test_validation_node(tool_schema: Any, use_message_key: bool):
validation_node = ValidationNode([tool_schema])
tool_name = getattr(tool_schema, "name", getattr(tool_schema, "__name__", None))
inputs = [
AIMessage(
"hi?",
tool_calls=[
{
"name": tool_name,
"args": {"some_val": 1, "some_other_val": "foo"},
"id": "some 0",
},
{
"name": tool_name,
# Wrong type for some_val
"args": {"some_val": "bar", "some_other_val": "foo"},
"id": "some 1",
},
],
),
]
if use_message_key:
inputs = {"messages": inputs}
result = await validation_node.ainvoke(inputs)
if use_message_key:
result = result["messages"]
def check_results(messages: list):
assert len(messages) == 2
assert all(m.type == "tool" for m in messages)
assert not messages[0].additional_kwargs.get("is_error")
assert messages[1].additional_kwargs.get("is_error")
check_results(result)
result_sync = validation_node.invoke(inputs)
if use_message_key:
result_sync = result_sync["messages"]
check_results(result_sync)