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

chore(prebuilt): rework structured outputs -- type safety, etc (#5962)

Browse Source 
* make `_SchemaSpec` private
* Add ability to customize message used in artificial tool response
This commit is contained in:
Sydney Runkle
2025-08-20 08:52:51 -04:00
committed by GitHub
parent 4151861ca2
commit e670815780
5 changed files with 117 additions and 143 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py
+31 -26
View File
@@ -5,10 +5,10 @@ from typing import (
Any,
Awaitable,
Callable,
Generic,
Optional,
Sequence,
Type,
TypeVar,
Union,
cast,
get_type_hints,
@@ -33,7 +33,7 @@ from langchain_core.runnables import (
)
from langchain_core.tools import BaseTool
from pydantic import BaseModel
from typing_extensions import Annotated, NotRequired, TypedDict
from typing_extensions import Annotated, NotRequired, TypedDict, TypeVar
from langgraph._internal._runnable import RunnableCallable, RunnableLike
from langgraph._internal._typing import MISSING
@@ -48,7 +48,6 @@ from langgraph.prebuilt._internal._typing import (
from langgraph.prebuilt.responses import (
OutputToolBinding,
ResponseFormat,
SchemaSpec,
ToolOutput,
)
from langgraph.prebuilt.tool_node import ToolNode
@@ -58,11 +57,12 @@ from langgraph.types import Checkpointer, Command, Send
from langgraph.typing import ContextT
from langgraph.warnings import LangGraphDeprecatedSinceV10
StructuredResponse = Union[dict, BaseModel]
StructuredResponseSchema = Union[dict, type[BaseModel]]
F = TypeVar("F", bound=Callable[..., Any])
StructuredResponseT = TypeVar(
"StructuredResponseT", bound=Union[dict, BaseModel, None], default=None
)
# We create the AgentState that we will pass around
# This simply involves a list of messages
@@ -84,16 +84,18 @@ class AgentStatePydantic(BaseModel):
remaining_steps: RemainingSteps = 25
class AgentStateWithStructuredResponse(AgentState):
class AgentStateWithStructuredResponse(AgentState, Generic[StructuredResponseT]):
"""The state of the agent with a structured response."""
structured_response: StructuredResponse
structured_response: StructuredResponseT
class AgentStateWithStructuredResponsePydantic(AgentStatePydantic):
class AgentStateWithStructuredResponsePydantic(
AgentStatePydantic, Generic[StructuredResponseT]
):
"""The state of the agent with a structured response."""
structured_response: StructuredResponse
structured_response: StructuredResponseT
StateSchema = TypeVar("StateSchema", bound=Union[AgentState, AgentStatePydantic])
@@ -185,7 +187,7 @@ def _validate_chat_history(
raise ValueError(error_message)
class _AgentBuilder:
class _AgentBuilder(Generic[StructuredResponseT]):
"""Internal builder class for constructing and agent."""
def __init__(
@@ -198,7 +200,7 @@ class _AgentBuilder:
tools: Union[Sequence[Union[BaseTool, Callable, dict[str, Any]]], ToolNode],
*,
prompt: Optional[Prompt] = None,
response_format: Optional[ResponseFormat] = None,
response_format: Optional[ResponseFormat[StructuredResponseT]] = None,
pre_model_hook: Optional[RunnableLike] = None,
post_model_hook: Optional[RunnableLike] = None,
state_schema: Optional[StateSchemaType] = None,
@@ -256,15 +258,15 @@ class _AgentBuilder:
Future strategies (json_mode, guided) will have separate setup methods.
"""
self.structured_output_tools: dict[str, OutputToolBinding] = {}
self.structured_output_tools: dict[
str, OutputToolBinding[StructuredResponseT]
] = {}
if self.response_format is not None:
response_format = self.response_format
# Handle UsingToolStrategy wrapper
if isinstance(response_format, ToolOutput):
# Use tools strategy - process each ResponseSchema in the UsingToolStrategy
for response_schema in response_format.schemas:
# Use the factory method to create OutputToolBinding
# check if response_format.schema is a union
for response_schema in response_format.schema_specs:
structured_tool_info = OutputToolBinding.from_schema_spec(
response_schema
)
@@ -275,7 +277,7 @@ class _AgentBuilder:
# This shouldn't happen with the new ResponseFormat type, but keeping for safety
raise ValueError(
f"Unsupported response_format type: {type(response_format)}. "
f"Expected UsingToolStrategy."
f"Expected ToolOutput."
)
def _setup_state_schema(self) -> None:
@@ -328,12 +330,17 @@ class _AgentBuilder:
"Behavior has not yet been defined in this case."
)
if isinstance(self.response_format, ToolOutput):
tool_message_content = self.response_format.tool_message_content
else:
tool_message_content = "ok!"
if len(structured_tool_calls) == 1:
tool_call = structured_tool_calls[0]
messages = [
response,
ToolMessage(
content="ok!",
content=tool_message_content,
tool_call_id=tool_call["id"],
name=tool_call["name"],
),
@@ -383,10 +390,8 @@ class _AgentBuilder:
if len(all_tools) > 0:
# Check if we need to force tool use for structured output
tool_choice = None
if (
self.response_format is not None
and isinstance(self.response_format, ToolOutput)
and self.response_format.tool_choice == "required"
if self.response_format is not None and isinstance(
self.response_format, ToolOutput
):
tool_choice = "any"
@@ -778,7 +783,7 @@ def create_react_agent(
tools: Union[Sequence[Union[BaseTool, Callable, dict[str, Any]]], ToolNode],
*,
prompt: Optional[Prompt] = None,
response_format: Optional[Union[ToolOutput, StructuredResponseSchema]] = None,
response_format: Optional[Union[ToolOutput, type[StructuredResponseT]]] = None,
pre_model_hook: Optional[RunnableLike] = None,
post_model_hook: Optional[RunnableLike] = None,
state_schema: Optional[StateSchemaType] = None,
@@ -985,9 +990,9 @@ def create_react_agent(
# Then it's a pydantic model or JSONSchema. We'll automatically convert
# it to the tool output strategy as it is widely supported.
response_format = ToolOutput(
schemas=[SchemaSpec(response_format)],
tool_choice="required",
schema=response_format,
)
elif isinstance(response_format, tuple):
if len(response_format) == 2:
raise ValueError(
libs/prebuilt/langgraph/prebuilt/responses.py
+45 -37
View File
@@ -2,33 +2,39 @@
from __future__ import annotations
import sys
from dataclasses import dataclass
from typing import Any, Generic, Literal, Optional, Sequence, Type, TypeVar
from typing import Any, Generic, Literal, TypeVar, Union, cast, get_args, get_origin
from langchain_core.tools import BaseTool
from langchain_core.tools import tool as create_tool
from pydantic import BaseModel
from typing_extensions import Self
# For now, we support only Pydantic models as schemas.
SchemaT = TypeVar("SchemaT", bound=type[BaseModel])
Schema = TypeVar("Schema", bound=BaseModel)
SchemaT = TypeVar("SchemaT")
ToolChoice = Literal["required", "auto"]
"""Required: model must call a tool; auto: model may respond with free text."""
if sys.version_info >= (3, 10):
from types import UnionType
else:
UnionType = Union
@dataclass(init=False)
class SchemaSpec(Generic[SchemaT]):
class _SchemaSpec(Generic[SchemaT]):
"""Describes a structured output schema."""
schema: SchemaT
schema: type[SchemaT]
"""The schema for the response, can be a dict or a Pydantic model."""
name: Optional[str] = None
"""Name of the schema, used for tool calling.
name: str | None = None
"""Name of the schema, used for tool calling.
If not provided, the name will be the model name.
"""
description: Optional[str] = None
description: str | None = None
"""Custom description of the schema.
If not provided, provided will use the model's docstring.
@@ -38,10 +44,10 @@ class SchemaSpec(Generic[SchemaT]):
def __init__(
self,
schema: SchemaT,
schema: type[SchemaT],
*,
name: Optional[str] = None,
description: Optional[str] = None,
name: str | None = None,
description: str | None = None,
strict: bool = False,
) -> None:
"""Initialize SchemaSpec with schema and optional parameters."""
@@ -52,31 +58,33 @@ class SchemaSpec(Generic[SchemaT]):
@dataclass(init=False)
class ToolOutput:
class ToolOutput(Generic[SchemaT]):
"""Use a tool calling strategy for model responses."""
schemas: Sequence[SchemaSpec]
"""Schemas for the tool calls."""
tool_choice: ToolChoice
"""Whether to require tool calling or allow the model to choose.
- "required": The model must use tool calling.
- "auto": The model can choose whether to use tool calling.
Use `auto` if you want the agent to be able to respond with a non structured
response to ask clarifying questions.
"""
schema: type[SchemaT]
"""Schema for the tool calls."""
tool_message_content: str
"""The content of the tool message to be returned when the model calls an artificial structured output tool."""
schema_specs: list[_SchemaSpec[SchemaT]]
"""Schema specs for the tool calls."""
def __init__(
self, schemas: Sequence[SchemaSpec], *, tool_choice: ToolChoice = "required"
self, schema: type[SchemaT], tool_message_content: str = "ok!"
) -> None:
"""Initialize ToolOutput with schemas and tool choice."""
self.schemas = schemas
self.tool_choice = tool_choice
"""Initialize ToolOutput with schemas and tool message content."""
self.schema = schema
self.tool_message_content = tool_message_content
if get_origin(schema) in (UnionType, Union):
self.schema_specs = [_SchemaSpec(s) for s in get_args(schema)]
else:
self.schema_specs = [_SchemaSpec(schema)]
@dataclass
class OutputToolBinding(Generic[Schema]):
class OutputToolBinding(Generic[SchemaT]):
"""Information for tracking structured output tool metadata.
This contains all necessary information to handle structured responses
@@ -84,17 +92,17 @@ class OutputToolBinding(Generic[Schema]):
and the corresponding tool implementation used by the tools strategy.
"""
schema: Type[Schema]
schema: type[SchemaT]
"""The original schema provided for structured output (Pydantic model or dict schema)."""
schema_kind: Literal["pydantic"]
"""Classification of the schema type for proper response construction."""
tool: BaseTool
"""LangChain tool instance created from the schema for model binding."""
@classmethod
def from_schema_spec(
cls, schema_spec: SchemaSpec[Type[Schema]]
) -> "OutputToolBinding":
def from_schema_spec(cls, schema_spec: _SchemaSpec[SchemaT]) -> Self:
"""Create an OutputToolBinding instance from a SchemaSpec.
Args:
@@ -129,12 +137,12 @@ class OutputToolBinding(Generic[Schema]):
tool = tool_creator(schema, **kwargs)
return cls(
schema=schema,
schema=cast(type[SchemaT], schema),
schema_kind=schema_kind,
tool=tool,
)
def parse(self, tool_args: dict[str, Any]) -> Schema:
def parse(self, tool_args: dict[str, Any]) -> SchemaT:
"""Parse tool arguments according to the schema.
Args:
@@ -164,4 +172,4 @@ class OutputToolBinding(Generic[Schema]):
# TODO: Add support for built-in structured responses (e.g., openai, grok)
ResponseFormat = ToolOutput
ResponseFormat = ToolOutput[SchemaT]
libs/prebuilt/tests/model.py
+5 -4
View File
@@ -2,6 +2,7 @@ from typing import (
Any,
Callable,
Dict,
Generic,
List,
Literal,
Optional,
@@ -22,12 +23,12 @@ from langchain_core.runnables import Runnable, RunnableLambda
from langchain_core.tools import BaseTool
from pydantic import BaseModel
from langgraph.prebuilt.chat_agent_executor import StructuredResponse
from langgraph.prebuilt.chat_agent_executor import StructuredResponseT
class FakeToolCallingModel(BaseChatModel):
class FakeToolCallingModel(BaseChatModel, Generic[StructuredResponseT]):
tool_calls: Optional[list[list[ToolCall]]] = None
structured_response: Optional[StructuredResponse] = None
structured_response: Optional[StructuredResponseT] = None
index: int = 0
tool_style: Literal["openai", "anthropic"] = "openai"
@@ -57,7 +58,7 @@ class FakeToolCallingModel(BaseChatModel):
def with_structured_output(
self, schema: Type[BaseModel]
) -> Runnable[LanguageModelInput, StructuredResponse]:
) -> Runnable[LanguageModelInput, StructuredResponseT]:
if self.structured_response is None:
raise ValueError("Structured response is not set")
libs/prebuilt/tests/test_react_agent.py
+1 -1
View File
@@ -433,7 +433,7 @@ def test_react_agent_with_structured_response() -> None:
return "The weather is sunny and 75°F."
expected_structured_response = WeatherResponse(temperature=75)
model = FakeToolCallingModel(
model = FakeToolCallingModel[WeatherResponse](
tool_calls=tool_calls, structured_response=expected_structured_response
)
agent = create_react_agent(
libs/prebuilt/tests/test_responses.py
+35 -75
View File
@@ -1,17 +1,18 @@
"""Unit tests for langgraph.prebuilt.responses module."""
from typing import Union
import pytest
from pydantic import BaseModel
from langgraph.prebuilt.responses import (
OutputToolBinding,
ResponseFormat,
SchemaSpec,
ToolOutput,
_SchemaSpec,
)
class TestModel(BaseModel):
class _TestModel(BaseModel):
"""A test model for structured output."""
name: str
@@ -31,54 +32,29 @@ class EmptyDocModel(BaseModel):
data: str
class TestSchemaSpec:
"""Test SchemaSpec dataclass."""
def test_basic_creation(self):
"""Test basic SchemaSpec creation."""
schema = SchemaSpec(schema=TestModel)
assert schema.schema == TestModel
assert schema.name is None
assert schema.description is None
assert schema.strict is False
def test_creation_with_all_fields(self):
"""Test SchemaSpec creation with all fields."""
schema = SchemaSpec(
schema=TestModel,
name="custom_test_model",
description="A custom description",
strict=True,
)
assert schema.schema == TestModel
assert schema.name == "custom_test_model"
assert schema.description == "A custom description"
assert schema.strict is True
class TestUsingToolStrategy:
"""Test UsingToolStrategy dataclass."""
def test_basic_creation(self):
"""Test basic UsingToolStrategy creation."""
schema = SchemaSpec(schema=TestModel)
strategy = ToolOutput(schemas=[schema])
assert len(strategy.schemas) == 1
assert strategy.schemas[0] == schema
assert strategy.tool_choice == "required" # default
def test_creation_with_auto_tool_choice(self):
"""Test UsingToolStrategy creation with auto tool choice."""
schema = SchemaSpec(schema=TestModel)
strategy = ToolOutput(schemas=[schema], tool_choice="auto")
assert strategy.tool_choice == "auto"
strategy = ToolOutput(schema=_TestModel)
assert strategy.schema == _TestModel
assert strategy.tool_message_content == "ok!"
assert len(strategy.schema_specs) == 1
def test_multiple_schemas(self):
"""Test UsingToolStrategy with multiple schemas."""
schema1 = SchemaSpec(schema=TestModel)
schema2 = SchemaSpec(schema=CustomModel)
strategy = ToolOutput(schemas=[schema1, schema2])
assert len(strategy.schemas) == 2
strategy = ToolOutput(schema=Union[_TestModel, CustomModel])
assert len(strategy.schema_specs) == 2
assert strategy.schema_specs[0].schema == _TestModel
assert strategy.schema_specs[1].schema == CustomModel
def test_schema_with_tool_message_content(self):
"""Test UsingToolStrategy with tool message content."""
strategy = ToolOutput(schema=_TestModel, tool_message_content="custom message")
assert strategy.schema == _TestModel
assert strategy.tool_message_content == "custom message"
assert len(strategy.schema_specs) == 1
class TestOutputToolBinding:
@@ -86,24 +62,24 @@ class TestOutputToolBinding:
def test_from_schema_spec_basic(self):
"""Test basic OutputToolBinding creation from SchemaSpec."""
schema_spec = SchemaSpec(schema=TestModel)
schema_spec = _SchemaSpec(schema=_TestModel)
tool_binding = OutputToolBinding.from_schema_spec(schema_spec)
assert tool_binding.schema == TestModel
assert tool_binding.schema == _TestModel
assert tool_binding.schema_kind == "pydantic"
assert tool_binding.tool is not None
assert tool_binding.tool.name == "TestModel"
assert tool_binding.tool.name == "_TestModel"
def test_from_schema_spec_with_custom_name(self):
"""Test OutputToolBinding creation with custom name."""
schema_spec = SchemaSpec(schema=TestModel, name="custom_tool_name")
schema_spec = _SchemaSpec(schema=_TestModel, name="custom_tool_name")
tool_binding = OutputToolBinding.from_schema_spec(schema_spec)
assert tool_binding.tool.name == "custom_tool_name"
def test_from_schema_spec_with_custom_description(self):
"""Test OutputToolBinding creation with custom description."""
schema_spec = SchemaSpec(
schema=TestModel, description="Custom tool description"
schema_spec = _SchemaSpec(
schema=_TestModel, description="Custom tool description"
)
tool_binding = OutputToolBinding.from_schema_spec(schema_spec)
@@ -111,7 +87,7 @@ class TestOutputToolBinding:
def test_from_schema_spec_with_model_docstring(self):
"""Test OutputToolBinding creation using model docstring as description."""
schema_spec = SchemaSpec(schema=CustomModel)
schema_spec = _SchemaSpec(schema=CustomModel)
tool_binding = OutputToolBinding.from_schema_spec(schema_spec)
assert tool_binding.tool.description == "Custom model with a custom docstring."
@@ -127,7 +103,7 @@ class TestOutputToolBinding:
# This should have the same docstring as BaseModel
pass
schema_spec = SchemaSpec(schema=DefaultDocModel)
schema_spec = _SchemaSpec(schema=DefaultDocModel)
tool_binding = OutputToolBinding.from_schema_spec(schema_spec)
# Should use empty description when model has default BaseModel docstring
@@ -135,26 +111,26 @@ class TestOutputToolBinding:
def test_parse_payload_pydantic_success(self):
"""Test successful parsing for Pydantic model."""
schema_spec = SchemaSpec(schema=TestModel)
schema_spec = _SchemaSpec(schema=_TestModel)
tool_binding = OutputToolBinding.from_schema_spec(schema_spec)
tool_args = {"name": "John", "age": 30}
result = tool_binding.parse(tool_args)
assert isinstance(result, TestModel)
assert isinstance(result, _TestModel)
assert result.name == "John"
assert result.age == 30
assert result.email == "default@example.com" # default value
def test_parse_payload_pydantic_validation_error(self):
"""Test parsing failure for invalid Pydantic data."""
schema_spec = SchemaSpec(schema=TestModel)
schema_spec = _SchemaSpec(schema=_TestModel)
tool_binding = OutputToolBinding.from_schema_spec(schema_spec)
# Missing required field 'name'
tool_args = {"age": 30}
with pytest.raises(ValueError, match="Failed to parse tool args to TestModel"):
with pytest.raises(ValueError, match="Failed to parse tool args to _TestModel"):
tool_binding.parse(tool_args)
def test_parse_payload_invalid_kind(self):
@@ -164,7 +140,7 @@ class TestOutputToolBinding:
mock_tool = Mock()
tool_binding = OutputToolBinding(
schema=TestModel,
schema=_TestModel,
schema_kind="invalid_kind", # type: ignore
tool=mock_tool,
)
@@ -189,36 +165,20 @@ class TestOutputToolBinding:
tool_binding.parse({"name": "test", "age": 25})
class TestResponseFormat:
"""Test ResponseFormat type alias."""
def test_response_format_is_using_tool_strategy(self):
"""Test that ResponseFormat is aliased to UsingToolStrategy."""
assert ResponseFormat is ToolOutput
def test_can_create_response_format(self):
"""Test that we can create ResponseFormat instances."""
schema = SchemaSpec(schema=TestModel)
response_format = ResponseFormat(schemas=[schema])
assert isinstance(response_format, ToolOutput)
assert len(response_format.schemas) == 1
class TestEdgeCases:
"""Test edge cases and error conditions."""
def test_empty_schemas_list(self) -> None:
"""Test UsingToolStrategy with empty schemas list."""
strategy = ToolOutput([SchemaSpec(EmptyDocModel)])
assert len(strategy.schemas) == 1
strategy = ToolOutput(EmptyDocModel)
assert len(strategy.schema_specs) == 1
@pytest.mark.skip(
reason="Need to fix bug in langchain-core for inheritance of doc-strings."
)
def test_base_model_doc_constant(self) -> None:
"""Test that BASE_MODEL_DOC constant is set correctly."""
binding = OutputToolBinding.from_schema_spec(SchemaSpec(EmptyDocModel))
binding = OutputToolBinding.from_schema_spec(_SchemaSpec(EmptyDocModel))
assert binding.tool.name == "EmptyDocModel"
assert (
binding.tool.description[:5] == ""