diff --git a/Makefile b/Makefile index fb74512f1..3b2ae1774 100644 --- a/Makefile +++ b/Makefile @@ -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 diff --git a/docs/docs/concepts/index.md b/docs/docs/concepts/index.md index 5655e2c32..b6f84e18c 100644 --- a/docs/docs/concepts/index.md +++ b/docs/docs/concepts/index.md @@ -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: diff --git a/docs/docs/reference/graphs.md b/docs/docs/reference/graphs.md index 01cfdd874..724ac724e 100644 --- a/docs/docs/reference/graphs.md +++ b/docs/docs/reference/graphs.md @@ -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 diff --git a/docs/docs/reference/prebuilt.md b/docs/docs/reference/prebuilt.md index 20fef193b..fc276aea0 100644 --- a/docs/docs/reference/prebuilt.md +++ b/docs/docs/reference/prebuilt.md @@ -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 \ No newline at end of file +::: langgraph.prebuilt.tools_condition + + +## ValidationNode + +```python +from langgraph.prebuilt import ValidationNode +``` + +::: langgraph.prebuilt.ValidationNode \ No newline at end of file diff --git a/langgraph/prebuilt/__init__.py b/langgraph/prebuilt/__init__.py index 48f20d2af..b05e0e660 100644 --- a/langgraph/prebuilt/__init__.py +++ b/langgraph/prebuilt/__init__.py @@ -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", ] diff --git a/langgraph/prebuilt/tool_validator.py b/langgraph/prebuilt/tool_validator.py new file mode 100644 index 000000000..9172fd7ea --- /dev/null +++ b/langgraph/prebuilt/tool_validator.py @@ -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} diff --git a/tests/test_prebuilt.py b/tests/test_prebuilt.py index a7d02083d..dd2345bb8 100644 --- a/tests/test_prebuilt.py +++ b/tests/test_prebuilt.py @@ -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)