feat(prebuilt): native structured output support w/ all sorts of models (#5961)

* Adds support for `NativeOutput` via a new `NativeOutput` dataclass
* Adds support for structured output specification via the following
(pydantic models already supported)
  * dataclasses
  * typed dicts
  * json schemas 
* Adds mocking support to support native strategies with
`FakeToolCallingModel`
* Add new default tool message when `tool_message_content` not provided
* Smart "selection" of native vs tool output based on provider support,
necessitates profiles down the line
  
Considered questions
* do we want to enforce docstrings? -- decided on no for now
* do we want to enforce names (titles) on json schemas? -- decided no
for now, defaulting to `structured_output`
* do we want to validate that json schemas coming in are valid? --
decided no for now
* do we want to validate model results against a given json schema? we
validate against all other types (typed dict, dataclass, etc) w/
pydantic -- decided no for now

TODO in future PRs:
* Figure out retry policy
* Add standard testing (handed off to @casparb)
* Further privatize certain structures (like the bindings) -- this is
low prio

---------

Co-authored-by: Sydney Runkle <sydneymarierunkle@gmail.com>
Co-authored-by: Sydney Runkle <54324534+sydney-runkle@users.noreply.github.com>
This commit is contained in:
Caspar Broekhuizen
2025-08-21 10:43:50 -04:00
committed by GitHub
co-authored by Sydney Runkle Sydney Runkle
parent 1cd1373788
commit f4cdeea6ad
8 changed files with 1097 additions and 159 deletions
@@ -1,6 +1,7 @@
from __future__ import annotations
import inspect
from dataclasses import asdict, is_dataclass
from typing import (
Any,
Awaitable,
@@ -42,6 +43,8 @@ from langgraph.prebuilt._internal._typing import (
SyncOrAsync,
)
from langgraph.prebuilt.responses import (
NativeOutput,
NativeOutputBinding,
OutputToolBinding,
ResponseFormat,
ToolOutput,
@@ -53,7 +56,8 @@ from langgraph.types import Checkpointer, Command, Send
from langgraph.typing import ContextT, StateT
StructuredResponseT = TypeVar(
"StructuredResponseT", bound=Union[dict, BaseModel, None], default=None
"StructuredResponseT",
default=None,
)
@@ -194,7 +198,7 @@ class _AgentBuilder(Generic[StateT, ContextT, StructuredResponseT]):
self._setup_tools()
self._setup_state_schema()
self._setup_structured_output_tools()
self._setup_structured_output()
self._setup_model()
def _setup_tools(self) -> None:
@@ -215,19 +219,26 @@ class _AgentBuilder(Generic[StateT, ContextT, StructuredResponseT]):
}
self._tool_calling_enabled = len(self._tool_classes) > 0
def _setup_structured_output_tools(self) -> None:
"""Set up structured output tools tracking for the tools strategy.
def _setup_structured_output(self) -> None:
"""Set up structured output tracking for "tools" and "native" strategies.
This method implements the "tools" strategy for structured output by:
"tools" strategy for structured output:
1. Converting response format schemas to LangChain tools
2. Creating metadata for proper response reconstruction
3. Handling both Pydantic models and dict schemas
Future strategies (json_mode, guided) will have separate setup methods.
"native" strategy for structured output:
1. Capturing the schema reference for later parsing
2. Binding provider-native response_format kwargs at model bind time
3. Parsing provider-enforced structured output directly into the schema
"""
self.structured_output_tools: dict[
str, OutputToolBinding[StructuredResponseT]
] = {}
self.native_output_binding: NativeOutputBinding[StructuredResponseT] | None = (
None
)
if self.response_format is not None:
response_format = self.response_format
@@ -240,6 +251,11 @@ class _AgentBuilder(Generic[StateT, ContextT, StructuredResponseT]):
self.structured_output_tools[structured_tool_info.tool.name] = (
structured_tool_info
)
elif isinstance(response_format, NativeOutput):
# Use native strategy - create NativeOutputBinding for parsing
self.native_output_binding = NativeOutputBinding.from_schema_spec(
response_format.schema_spec
)
else:
# This shouldn't happen with the new ResponseFormat type, but keeping for safety
raise ValueError(
@@ -293,33 +309,67 @@ class _AgentBuilder(Generic[StateT, ContextT, StructuredResponseT]):
"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=tool_message_content,
tool_call_id=tool_call["id"],
name=tool_call["name"],
),
]
structured_tool_binding = self.structured_output_tools[tool_call["name"]]
structured_response = structured_tool_binding.parse(tool_call["args"])
if isinstance(structured_response, BaseModel):
structured_response_dict = structured_response.model_dump()
elif is_dataclass(structured_response):
structured_response_dict = asdict(structured_response) # type: ignore[arg-type]
else:
structured_response_dict = cast(dict, structured_response)
tool_message_content = (
self.response_format.tool_message_content
if isinstance(self.response_format, ToolOutput)
and self.response_format.tool_message_content
else f"Returning structured response: {structured_response_dict}"
)
return Command(
update={
"messages": messages,
"messages": [
response,
ToolMessage(
content=tool_message_content,
tool_call_id=tool_call["id"],
name=tool_call["name"],
),
],
"structured_response": structured_response,
}
)
return None
def _apply_native_output_binding(
self, model: LanguageModelLike
) -> LanguageModelLike:
"""If native output is configured, bind provider-native kwargs onto the model."""
if not isinstance(self.response_format, NativeOutput):
return model
kwargs = self.response_format.to_model_kwargs()
model_with_native_output = model.bind(**kwargs)
return model_with_native_output
def _handle_structured_response_native(
self, response: AIMessage
) -> Optional[Command]:
"""If native output is configured and there are no tool calls, parse using NativeOutputBinding."""
if self.native_output_binding is None:
return None
if response.tool_calls:
# if the model chooses to call tools, we let the normal flow handle it
return None
structured_response = self.native_output_binding.parse(response)
return Command(
update={"messages": [response], "structured_response": structured_response}
)
def _setup_model(self) -> None:
"""Setup model-related attributes."""
self._is_dynamic_model = not isinstance(
@@ -363,9 +413,19 @@ class _AgentBuilder(Generic[StateT, ContextT, StructuredResponseT]):
all_tools, tool_choice=tool_choice
)
else:
model = cast(BaseChatModel, model).bind_tools(all_tools) # type: ignore[assignment]
# If native output is configured, bind tools with strict=True. Required for OpenAI.
if isinstance(self.response_format, NativeOutput):
model = cast(BaseChatModel, model).bind_tools( # type: ignore[assignment]
all_tools, strict=True
)
else:
model = cast(BaseChatModel, model).bind_tools(all_tools) # type: ignore[assignment]
# bind native structured-output kwargs
model = self._apply_native_output_binding(model) # type: ignore[arg-type]
# Extract just the model part for direct invocation
self._static_model: Optional[Runnable] = model # type: ignore[assignment]
self._static_model: Optional[Runnable] = model
else:
self._static_model = None
@@ -374,7 +434,8 @@ class _AgentBuilder(Generic[StateT, ContextT, StructuredResponseT]):
) -> LanguageModelLike:
"""Resolve the model to use, handling both static and dynamic models."""
if self._is_dynamic_model:
return self.model(state, runtime) # type: ignore[operator, arg-type]
dynamic_model = self.model(state, runtime) # type: ignore[operator, arg-type]
return self._apply_native_output_binding(dynamic_model) # type: ignore[arg-type]
else:
return self._static_model
@@ -390,7 +451,8 @@ class _AgentBuilder(Generic[StateT, ContextT, StructuredResponseT]):
resolved_model = await dynamic_model(state, runtime)
return resolved_model
elif self._is_dynamic_model:
return self.model(state, runtime) # type: ignore[arg-type, operator]
dynamic_model = self.model(state, runtime) # type: ignore[arg-type, assignment, operator]
return self._apply_native_output_binding(dynamic_model) # type: ignore[arg-type]
else:
return self._static_model
@@ -472,6 +534,11 @@ class _AgentBuilder(Generic[StateT, ContextT, StructuredResponseT]):
if structured_command:
return structured_command
# Native structured output
native_command = self._handle_structured_response_native(response)
if native_command:
return native_command
return {"messages": [response]}
async def acall_model(
@@ -507,6 +574,11 @@ class _AgentBuilder(Generic[StateT, ContextT, StructuredResponseT]):
if structured_command:
return structured_command
# Native structured output
native_command = self._handle_structured_response_native(response)
if native_command:
return native_command
return {"messages": [response]}
return RunnableCallable(call_model, acall_model)
@@ -703,6 +775,32 @@ class _AgentBuilder(Generic[StateT, ContextT, StructuredResponseT]):
return workflow
def _supports_native_structured_output(
model: Union[
str, BaseChatModel, SyncOrAsync[[StateT, Runtime[ContextT]], BaseChatModel]
],
) -> bool:
"""Check if a model supports native structured output.
TODO: replace with more robust model profiles.
"""
model_name: str | None = None
if isinstance(model, str):
model_name = model
elif isinstance(model, BaseChatModel):
model_name = getattr(model, "model_name", None)
return (
"grok" in model_name.lower()
or any(
part in model_name
for part in ["gpt-5", "gpt-4.1", "gpt-oss", "o3-pro", "o3-mini"]
)
if model_name
else False
)
def create_agent(
model: Union[
str,
@@ -713,7 +811,11 @@ def create_agent(
*,
prompt: Optional[Prompt] = None,
response_format: Optional[
Union[ToolOutput[StructuredResponseT], type[StructuredResponseT]]
Union[
ToolOutput[StructuredResponseT],
NativeOutput[StructuredResponseT],
type[StructuredResponseT],
]
] = None,
pre_model_hook: Optional[RunnableLike] = None,
post_model_hook: Optional[RunnableLike] = None,
@@ -888,28 +990,29 @@ def create_agent(
print(chunk)
```
"""
if response_format and not isinstance(response_format, ToolOutput):
# 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(
schema=response_format,
)
if response_format and not isinstance(response_format, (ToolOutput, NativeOutput)):
if _supports_native_structured_output(model):
response_format = NativeOutput(
schema=response_format,
)
else:
response_format = ToolOutput(
schema=response_format,
)
elif isinstance(response_format, tuple):
if len(response_format) == 2:
raise ValueError(
"Passing a 2-tuple as response_format is no longer supported. "
)
else:
# Can only be a ToolOutput or None at this point.
response_format = cast(Optional[ToolOutput], response_format)
# Create and configure the agent builder
builder = _AgentBuilder[StateT, ContextT, StructuredResponseT](
model=model,
tools=tools,
prompt=prompt,
response_format=response_format,
response_format=cast(
Union[ResponseFormat[StructuredResponseT], None], response_format
),
pre_model_hook=pre_model_hook,
post_model_hook=post_model_hook,
state_schema=state_schema,
+212 -68
View File
@@ -3,15 +3,15 @@
from __future__ import annotations
import sys
from dataclasses import dataclass
from typing import Any, Generic, Literal, TypeVar, Union, cast, get_args, get_origin
from dataclasses import dataclass, is_dataclass
from typing import Any, Generic, Literal, TypeVar, Union, 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
from langchain_core.messages import AIMessage
from langchain_core.tools import BaseTool, StructuredTool
from pydantic import BaseModel, TypeAdapter
from typing_extensions import Self, is_typeddict
# For now, we support only Pydantic models as schemas.
# Supported schema types: Pydantic models, dataclasses, TypedDict, JSON schema dicts
SchemaT = TypeVar("SchemaT")
@@ -20,31 +20,66 @@ if sys.version_info >= (3, 10):
else:
UnionType = Union
SchemaKind = Literal["pydantic", "dataclass", "typeddict", "json_schema"]
def _parse_with_schema(
schema: Union[type[SchemaT], dict], schema_kind: SchemaKind, data: dict[str, Any]
) -> Any:
"""Parse data using for any supported schema type.
Args:
schema: The schema type (Pydantic model, dataclass, or TypedDict)
data: The data to parse
Returns:
The parsed instance according to the schema type
Raises:
ValueError: If parsing fails
"""
if schema_kind == "json_schema":
return data
else:
try:
adapter: TypeAdapter[SchemaT] = TypeAdapter(schema)
return adapter.validate_python(data)
except Exception as e:
schema_name = getattr(schema, "__name__", str(schema))
raise ValueError(f"Failed to parse data to {schema_name}: {e}") from e
@dataclass(init=False)
class _SchemaSpec(Generic[SchemaT]):
"""Describes a structured output schema."""
schema: type[SchemaT]
"""The schema for the response, can be a dict or a Pydantic model."""
schema: Union[type[SchemaT], dict[str, Any]]
"""The schema for the response, can be a Pydantic model, dataclass, TypedDict, or JSON schema dict."""
name: str | None = None
name: str
"""Name of the schema, used for tool calling.
If not provided, the name will be the model name.
If not provided, the name will be the model name or "structured_output" if it's a JSON schema.
"""
description: str | None = None
description: str
"""Custom description of the schema.
If not provided, provided will use the model's docstring.
"""
schema_kind: SchemaKind
"""The kind of schema."""
json_schema: dict[str, Any]
"""JSON schema associated with the schema."""
strict: bool = False
"""Whether to enforce strict validation of the schema."""
def __init__(
self,
schema: type[SchemaT],
schema: Union[type[SchemaT], dict[str, Any]],
*,
name: str | None = None,
description: str | None = None,
@@ -52,26 +87,57 @@ class _SchemaSpec(Generic[SchemaT]):
) -> None:
"""Initialize SchemaSpec with schema and optional parameters."""
self.schema = schema
self.name = name
self.description = description
self.name = name or (
schema.get("title", "structured_output")
if isinstance(schema, dict)
else getattr(schema, "__name__", "structured_output")
)
self.description = description or (
schema.get("description", "")
if isinstance(schema, dict)
else getattr(schema, "__doc__", None) or ""
)
self.strict = strict
if isinstance(schema, dict):
self.schema_kind = "json_schema"
self.json_schema = schema
elif isinstance(schema, type) and issubclass(schema, BaseModel):
self.schema_kind = "pydantic"
self.json_schema = schema.model_json_schema()
elif is_dataclass(schema):
self.schema_kind = "dataclass"
self.json_schema = TypeAdapter(schema).json_schema()
elif is_typeddict(schema):
self.schema_kind = "typeddict"
self.json_schema = TypeAdapter(schema).json_schema()
else:
raise ValueError(
f"Unsupported schema type: {type(schema)}. "
f"Supported types: Pydantic models, dataclasses, TypedDicts, and JSON schema dicts."
)
@dataclass(init=False)
class ToolOutput(Generic[SchemaT]):
"""Use a tool calling strategy for model responses."""
schema: type[SchemaT]
schema: Union[type[SchemaT], dict[str, Any]]
"""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."""
tool_message_content: str | None
"""The content of the tool message to be returned when the model calls an artificial structured output tool."""
def __init__(
self, schema: type[SchemaT], tool_message_content: str = "ok!"
self,
schema: Union[type[SchemaT], dict[str, Any]],
tool_message_content: str | None = None,
) -> None:
"""Initialize ToolOutput with schemas and tool message content."""
self.schema = schema
@@ -83,6 +149,36 @@ class ToolOutput(Generic[SchemaT]):
self.schema_specs = [_SchemaSpec(schema)]
@dataclass(init=False)
class NativeOutput(Generic[SchemaT]):
"""Use the model provider's native structured output method."""
schema: Union[type[SchemaT], dict[str, Any]]
"""Schema for native mode."""
schema_spec: _SchemaSpec[SchemaT]
"""Schema spec for native mode."""
def __init__(
self,
schema: Union[type[SchemaT], dict[str, Any]],
) -> None:
self.schema = schema
self.schema_spec = _SchemaSpec(schema)
def to_model_kwargs(self) -> dict[str, Any]:
# OpenAI:
# - see https://platform.openai.com/docs/guides/structured-outputs
response_format = {
"type": "json_schema",
"json_schema": {
"name": self.schema_spec.name,
"schema": self.schema_spec.json_schema,
},
}
return {"response_format": response_format}
@dataclass
class OutputToolBinding(Generic[SchemaT]):
"""Information for tracking structured output tool metadata.
@@ -92,10 +188,10 @@ class OutputToolBinding(Generic[SchemaT]):
and the corresponding tool implementation used by the tools strategy.
"""
schema: type[SchemaT]
"""The original schema provided for structured output (Pydantic model or dict schema)."""
schema: Union[type[SchemaT], dict[str, Any]]
"""The original schema provided for structured output (Pydantic model, dataclass, TypedDict, or JSON schema dict)."""
schema_kind: Literal["pydantic"]
schema_kind: SchemaKind
"""Classification of the schema type for proper response construction."""
tool: BaseTool
@@ -111,35 +207,14 @@ class OutputToolBinding(Generic[SchemaT]):
Returns:
An OutputToolBinding instance with the appropriate tool created
"""
# Extract the actual schema from SchemaSpec
schema = schema_spec.schema
kwargs = {}
schema_kind: Literal["pydantic"]
# Use custom description if provided
if schema_spec.description:
kwargs["description"] = schema_spec.description
# Determine schema kind
if isinstance(schema, type) and issubclass(schema, BaseModel):
schema_kind = "pydantic"
else:
raise ValueError(
f"Unsupported schema type: {type(schema)}. "
f"Only Pydantic models are supported."
)
if schema_spec.name is not None:
tool_creator = create_tool(schema_spec.name)
else:
tool_creator = create_tool # type: ignore[assignment]
tool = tool_creator(schema, **kwargs)
return cls(
schema=cast(type[SchemaT], schema),
schema_kind=schema_kind,
tool=tool,
schema=schema_spec.schema,
schema_kind=schema_spec.schema_kind,
tool=StructuredTool(
args_schema=schema_spec.json_schema,
name=schema_spec.name,
description=schema_spec.description,
),
)
def parse(self, tool_args: dict[str, Any]) -> SchemaT:
@@ -154,22 +229,91 @@ class OutputToolBinding(Generic[SchemaT]):
Raises:
ValueError: If parsing fails
"""
if self.schema_kind == "pydantic":
if not isinstance(self.schema, type) or not issubclass(
self.schema, BaseModel
):
raise ValueError(
f"Expected Pydantic model class for 'pydantic' kind, got {type(self.schema)}"
)
try:
return self.schema(**tool_args)
except Exception as e:
raise ValueError(
f"Failed to parse tool args to {self.schema.__name__}: {e}"
) from e
else:
raise ValueError(f"Unsupported schema kind: {self.schema_kind}")
return _parse_with_schema(self.schema, self.schema_kind, tool_args)
# TODO: Add support for built-in structured responses (e.g., openai, grok)
ResponseFormat = ToolOutput[SchemaT]
@dataclass
class NativeOutputBinding(Generic[SchemaT]):
"""Information for tracking native structured output metadata.
This contains all necessary information to handle structured responses
generated via native provider output, including the original schema,
its type classification, and parsing logic for provider-enforced JSON.
"""
schema: Union[type[SchemaT], dict[str, Any]]
"""The original schema provided for structured output (Pydantic model, dataclass, TypedDict, or JSON schema dict)."""
schema_kind: SchemaKind
"""Classification of the schema type for proper response construction."""
@classmethod
def from_schema_spec(cls, schema_spec: _SchemaSpec[SchemaT]) -> Self:
"""Create a NativeOutputBinding instance from a SchemaSpec.
Args:
schema_spec: The SchemaSpec to convert
Returns:
A NativeOutputBinding instance for parsing native structured output
"""
return cls(
schema=schema_spec.schema,
schema_kind=schema_spec.schema_kind,
)
def parse(self, response: AIMessage) -> SchemaT:
"""Parse AIMessage content according to the schema.
Args:
response: The AI message containing the structured output
Returns:
The parsed response according to the schema
Raises:
ValueError: If text extraction, JSON parsing or schema validation fails
"""
# Extract text content from AIMessage and parse as JSON
raw_text = self._extract_text_content_from_message(response)
import json
try:
data = json.loads(raw_text)
except Exception as e:
schema_name = getattr(self.schema, "__name__", "structured_output")
raise ValueError(
f"Native structured output expected valid JSON for {schema_name}, but parsing failed: {e}."
) from e
# Parse according to schema
return _parse_with_schema(self.schema, self.schema_kind, data)
def _extract_text_content_from_message(self, message: AIMessage) -> str:
"""Extract text content from an AIMessage.
Args:
message: The AI message to extract text from
Returns:
The extracted text content
"""
content = message.content
if isinstance(content, str):
return content
if isinstance(content, list):
parts: list[str] = []
for c in content:
if isinstance(c, dict):
if c.get("type") == "text" and "text" in c:
parts.append(str(c["text"]))
elif "content" in c and isinstance(c["content"], str):
parts.append(c["content"])
else:
parts.append(str(c))
return "".join(parts)
return str(content)
ResponseFormat = Union[ToolOutput[SchemaT], NativeOutput[SchemaT]]
+36 -19
View File
@@ -1,3 +1,5 @@
import json
from dataclasses import asdict, is_dataclass
from typing import (
Any,
Callable,
@@ -19,7 +21,7 @@ from langchain_core.messages import (
ToolCall,
)
from langchain_core.outputs import ChatGeneration, ChatResult
from langchain_core.runnables import Runnable, RunnableLambda
from langchain_core.runnables import Runnable
from langchain_core.tools import BaseTool
from pydantic import BaseModel
@@ -27,7 +29,7 @@ from langgraph.prebuilt.chat_agent_executor import StructuredResponseT
class FakeToolCallingModel(BaseChatModel, Generic[StructuredResponseT]):
tool_calls: Optional[list[list[ToolCall]]] = None
tool_calls: Optional[Union[list[list[ToolCall]], list[list[dict]]]] = None
structured_response: Optional[StructuredResponseT] = None
index: int = 0
tool_style: Literal["openai", "anthropic"] = "openai"
@@ -40,15 +42,38 @@ class FakeToolCallingModel(BaseChatModel, Generic[StructuredResponseT]):
**kwargs: Any,
) -> ChatResult:
"""Top Level call"""
messages_string = "-".join([m.content for m in messages])
tool_calls = (
self.tool_calls[self.index % len(self.tool_calls)]
if self.tool_calls
else []
)
message = AIMessage(
content=messages_string, id=str(self.index), tool_calls=tool_calls.copy()
)
rf = kwargs.get("response_format")
is_native = isinstance(rf, dict) and rf.get("type") == "json_schema"
if is_native:
print("NATIVE. tool_calls: ", self.tool_calls)
if self.tool_calls:
if is_native:
tool_calls = (
self.tool_calls[self.index]
if self.index < len(self.tool_calls)
else []
)
else:
tool_calls = self.tool_calls[self.index % len(self.tool_calls)]
else:
tool_calls = []
if is_native and not tool_calls:
if isinstance(self.structured_response, BaseModel):
content_obj = self.structured_response.model_dump()
elif is_dataclass(self.structured_response):
content_obj = asdict(self.structured_response)
elif isinstance(self.structured_response, dict):
content_obj = self.structured_response
message = AIMessage(content=json.dumps(content_obj), id=str(self.index))
else:
messages_string = "-".join([m.content for m in messages])
message = AIMessage(
content=messages_string,
id=str(self.index),
tool_calls=tool_calls.copy(),
)
self.index += 1
return ChatResult(generations=[ChatGeneration(message=message)])
@@ -56,14 +81,6 @@ class FakeToolCallingModel(BaseChatModel, Generic[StructuredResponseT]):
def _llm_type(self) -> str:
return "fake-tool-call-model"
def with_structured_output(
self, schema: Type[BaseModel]
) -> Runnable[LanguageModelInput, StructuredResponseT]:
if self.structured_response is None:
raise ValueError("Structured response is not set")
return RunnableLambda(lambda x: self.structured_response)
def bind_tools(
self,
tools: Sequence[Union[Dict[str, Any], Type[BaseModel], Callable, BaseTool]],
@@ -0,0 +1,46 @@
[
{
"name": "updated structured response",
"responseFormat": [
{
"type": "object",
"properties": {
"name": { "type": "string" },
"role": { "type": "string" }
},
"required": ["name", "role"]
},
{
"type": "object",
"properties": {
"name": { "type": "string" },
"department": { "type": "string" }
},
"required": ["name", "department"]
}
],
"assertionsByInvocation": [
{
"prompt": "What is the role of Sabine?",
"toolsWithExpectedCalls": {
"getEmployeeRole": 1,
"getEmployeeDepartment": 0
},
"expectedLastMessage": "Returning structured response: {'name': 'Sabine', 'role': 'Developer'}",
"expectedStructuredResponse": { "name": "Sabine", "role": "Developer" },
"llmRequestCount": 2
},
{
"prompt": "In which department does Henrik work?",
"toolsWithExpectedCalls": {
"getEmployeeRole": 1,
"getEmployeeDepartment": 1
},
"expectedLastMessage": "Returning structured response: {'name': 'Henrik', 'department': 'IT'}",
"expectedStructuredResponse": { "name": "Henrik", "department": "IT" },
"llmRequestCount": 4
}
]
}
]
+3 -3
View File
@@ -447,7 +447,7 @@ def test_react_agent_with_structured_response() -> None:
"ai", # "What's the weather?"
"tool", # "The weather is sunny and 75°F."
"ai", # structured response
"tool", # ok!
"tool", # artificial tool message
]
assert [m.content for m in response["messages"]] == [
@@ -455,7 +455,7 @@ def test_react_agent_with_structured_response() -> None:
"What's the weather?",
"The weather is sunny and 75°F.",
"What's the weather?-What's the weather?-The weather is sunny and 75°F.",
"ok!",
"Returning structured response: {'temperature': 75.0}",
]
@@ -1608,7 +1608,7 @@ def test_post_model_hook_with_structured_output() -> None:
],
),
_AnyIdToolMessage(
content="ok!",
content="Returning structured response: {'temperature': 75.0}",
name="WeatherResponse",
tool_call_id="2",
),
+507
View File
@@ -0,0 +1,507 @@
"""Test suite for create_react_agent with structured output response_format permutations."""
from dataclasses import dataclass
from typing import Union
import pytest
from langchain_core.messages import HumanMessage
from pydantic import BaseModel, Field
from typing_extensions import TypedDict
from langgraph.prebuilt import create_agent
from langgraph.prebuilt.responses import NativeOutput, ToolOutput
from tests.model import FakeToolCallingModel
try:
from langchain_openai import ChatOpenAI
except ImportError:
skip_openai_integration_tests = True
else:
skip_openai_integration_tests = False
# Test data models
class WeatherBaseModel(BaseModel):
"""Weather response."""
temperature: float = Field(description="The temperature in fahrenheit")
condition: str = Field(description="Weather condition")
@dataclass
class WeatherDataclass:
"""Weather response."""
temperature: float
condition: str
class WeatherTypedDict(TypedDict):
"""Weather response."""
temperature: float
condition: str
weather_json_schema = {
"type": "object",
"properties": {
"temperature": {"type": "number", "description": "Temperature in fahrenheit"},
"condition": {"type": "string", "description": "Weather condition"},
},
"title": "weather_schema",
"required": ["temperature", "condition"],
}
class LocationResponse(BaseModel):
city: str = Field(description="The city name")
country: str = Field(description="The country name")
def get_weather() -> str:
"""Get the weather."""
return "The weather is sunny and 75°F."
def get_location() -> str:
"""Get the current location."""
return "You are in New York, USA."
# Standardized test data
WEATHER_DATA = {"temperature": 75.0, "condition": "sunny"}
LOCATION_DATA = {"city": "New York", "country": "USA"}
# Standardized expected responses
EXPECTED_WEATHER_PYDANTIC = WeatherBaseModel(**WEATHER_DATA)
EXPECTED_WEATHER_DATACLASS = WeatherDataclass(**WEATHER_DATA)
EXPECTED_WEATHER_DICT: WeatherTypedDict = {"temperature": 75.0, "condition": "sunny"}
EXPECTED_LOCATION = LocationResponse(**LOCATION_DATA)
class TestResponseFormatAsModel:
def test_pydantic_model(self) -> None:
"""Test response_format as Pydantic model."""
tool_calls = [
[{"args": {}, "id": "1", "name": "get_weather"}],
[
{
"name": "WeatherBaseModel",
"id": "2",
"args": WEATHER_DATA,
}
],
]
model = FakeToolCallingModel(tool_calls=tool_calls)
agent = create_agent(model, [get_weather], response_format=WeatherBaseModel)
response = agent.invoke({"messages": [HumanMessage("What's the weather?")]})
assert response["structured_response"] == EXPECTED_WEATHER_PYDANTIC
assert len(response["messages"]) == 5
def test_dataclass(self) -> None:
"""Test response_format as dataclass."""
tool_calls = [
[{"args": {}, "id": "1", "name": "get_weather"}],
[
{
"name": "WeatherDataclass",
"id": "2",
"args": WEATHER_DATA,
}
],
]
model = FakeToolCallingModel(tool_calls=tool_calls)
agent = create_agent(model, [get_weather], response_format=WeatherDataclass)
response = agent.invoke({"messages": [HumanMessage("What's the weather?")]})
assert response["structured_response"] == EXPECTED_WEATHER_DATACLASS
assert len(response["messages"]) == 5
def test_typed_dict(self) -> None:
"""Test response_format as TypedDict."""
tool_calls = [
[{"args": {}, "id": "1", "name": "get_weather"}],
[
{
"name": "WeatherTypedDict",
"id": "2",
"args": WEATHER_DATA,
}
],
]
model = FakeToolCallingModel(tool_calls=tool_calls)
agent = create_agent(model, [get_weather], response_format=WeatherTypedDict)
response = agent.invoke({"messages": [HumanMessage("What's the weather?")]})
assert response["structured_response"] == EXPECTED_WEATHER_DICT
assert len(response["messages"]) == 5
def test_json_schema(self) -> None:
"""Test response_format as JSON schema."""
tool_calls = [
[{"args": {}, "id": "1", "name": "get_weather"}],
[
{
"name": "weather_schema",
"id": "2",
"args": WEATHER_DATA,
}
],
]
model = FakeToolCallingModel(tool_calls=tool_calls)
agent = create_agent(model, [get_weather], response_format=weather_json_schema)
response = agent.invoke({"messages": [HumanMessage("What's the weather?")]})
assert response["structured_response"] == EXPECTED_WEATHER_DICT
assert len(response["messages"]) == 5
class TestResponseFormatAsToolOutput:
def test_pydantic_model(self) -> None:
"""Test response_format as ToolOutput with Pydantic model."""
tool_calls = [
[{"args": {}, "id": "1", "name": "get_weather"}],
[
{
"name": "WeatherBaseModel",
"id": "2",
"args": WEATHER_DATA,
}
],
]
model = FakeToolCallingModel(tool_calls=tool_calls)
agent = create_agent(
model, [get_weather], response_format=ToolOutput(WeatherBaseModel)
)
response = agent.invoke({"messages": [HumanMessage("What's the weather?")]})
assert response["structured_response"] == EXPECTED_WEATHER_PYDANTIC
assert len(response["messages"]) == 5
def test_dataclass(self) -> None:
"""Test response_format as ToolOutput with dataclass."""
tool_calls = [
[{"args": {}, "id": "1", "name": "get_weather"}],
[
{
"name": "WeatherDataclass",
"id": "2",
"args": WEATHER_DATA,
}
],
]
model = FakeToolCallingModel(tool_calls=tool_calls)
agent = create_agent(
model, [get_weather], response_format=ToolOutput(WeatherDataclass)
)
response = agent.invoke({"messages": [HumanMessage("What's the weather?")]})
assert response["structured_response"] == EXPECTED_WEATHER_DATACLASS
assert len(response["messages"]) == 5
def test_typed_dict(self) -> None:
"""Test response_format as ToolOutput with TypedDict."""
tool_calls = [
[{"args": {}, "id": "1", "name": "get_weather"}],
[
{
"name": "WeatherTypedDict",
"id": "2",
"args": WEATHER_DATA,
}
],
]
model = FakeToolCallingModel(tool_calls=tool_calls)
agent = create_agent(
model, [get_weather], response_format=ToolOutput(WeatherTypedDict)
)
response = agent.invoke({"messages": [HumanMessage("What's the weather?")]})
assert response["structured_response"] == EXPECTED_WEATHER_DICT
assert len(response["messages"]) == 5
def test_json_schema(self) -> None:
"""Test response_format as ToolOutput with JSON schema."""
tool_calls = [
[{"args": {}, "id": "1", "name": "get_weather"}],
[
{
"name": "weather_schema",
"id": "2",
"args": WEATHER_DATA,
}
],
]
model = FakeToolCallingModel(tool_calls=tool_calls)
agent = create_agent(
model, [get_weather], response_format=ToolOutput(weather_json_schema)
)
response = agent.invoke({"messages": [HumanMessage("What's the weather?")]})
assert response["structured_response"] == EXPECTED_WEATHER_DICT
assert len(response["messages"]) == 5
def test_union_of_types(self) -> None:
"""Test response_format as ToolOutput with Union of various types."""
# Test with WeatherBaseModel
tool_calls = [
[{"args": {}, "id": "1", "name": "get_weather"}],
[
{
"name": "WeatherBaseModel",
"id": "2",
"args": WEATHER_DATA,
}
],
]
model = FakeToolCallingModel[Union[WeatherBaseModel, LocationResponse]](
tool_calls=tool_calls
)
agent = create_agent(
model,
[get_weather, get_location],
response_format=ToolOutput(Union[WeatherBaseModel, LocationResponse]),
)
response = agent.invoke({"messages": [HumanMessage("What's the weather?")]})
assert response["structured_response"] == EXPECTED_WEATHER_PYDANTIC
assert len(response["messages"]) == 5
# Test with LocationResponse
tool_calls_location = [
[{"args": {}, "id": "1", "name": "get_location"}],
[
{
"name": "LocationResponse",
"id": "2",
"args": LOCATION_DATA,
}
],
]
model_location = FakeToolCallingModel(tool_calls=tool_calls_location)
agent_location = create_agent(
model_location,
[get_weather, get_location],
response_format=ToolOutput(Union[WeatherBaseModel, LocationResponse]),
)
response_location = agent_location.invoke(
{"messages": [HumanMessage("Where am I?")]}
)
assert response_location["structured_response"] == EXPECTED_LOCATION
assert len(response_location["messages"]) == 5
def test_multiple_tool_messages(self) -> None:
"""Test response_format as ToolOutput with Pydantic model."""
tool_calls = [
[{"args": {}, "id": "1", "name": "get_weather"}],
[
{
"name": "WeatherBaseModel",
"id": "2",
"args": WEATHER_DATA,
},
{
"name": "WeatherDataclass",
"id": "3",
"args": WEATHER_DATA,
},
],
]
model = FakeToolCallingModel(tool_calls=tool_calls)
agent = create_agent(
model,
[get_weather],
response_format=ToolOutput(Union[WeatherBaseModel, WeatherDataclass]),
)
with pytest.raises(
AssertionError,
match="Model incorrectly returned multiple structured responses.",
):
agent.invoke({"messages": [HumanMessage("What's the weather?")]})
class TestResponseFormatAsNativeOutput:
def test_pydantic_model(self) -> None:
"""Test response_format as NativeOutput with Pydantic model."""
tool_calls = [
[{"args": {}, "id": "1", "name": "get_weather"}],
]
model = FakeToolCallingModel[WeatherBaseModel](
tool_calls=tool_calls, structured_response=EXPECTED_WEATHER_PYDANTIC
)
agent = create_agent(
model, [get_weather], response_format=NativeOutput(WeatherBaseModel)
)
response = agent.invoke({"messages": [HumanMessage("What's the weather?")]})
assert response["structured_response"] == EXPECTED_WEATHER_PYDANTIC
assert len(response["messages"]) == 4
def test_dataclass(self) -> None:
"""Test response_format as NativeOutput with dataclass."""
tool_calls = [
[{"args": {}, "id": "1", "name": "get_weather"}],
]
model = FakeToolCallingModel[WeatherDataclass](
tool_calls=tool_calls, structured_response=EXPECTED_WEATHER_DATACLASS
)
agent = create_agent(
model, [get_weather], response_format=NativeOutput(WeatherDataclass)
)
response = agent.invoke(
{"messages": [HumanMessage("What's the weather?")]},
)
assert response["structured_response"] == EXPECTED_WEATHER_DATACLASS
assert len(response["messages"]) == 4
def test_typed_dict(self) -> None:
"""Test response_format as NativeOutput with TypedDict."""
tool_calls = [
[{"args": {}, "id": "1", "name": "get_weather"}],
]
model = FakeToolCallingModel[WeatherTypedDict](
tool_calls=tool_calls, structured_response=EXPECTED_WEATHER_DICT
)
agent = create_agent(
model, [get_weather], response_format=NativeOutput(WeatherTypedDict)
)
response = agent.invoke({"messages": [HumanMessage("What's the weather?")]})
assert response["structured_response"] == EXPECTED_WEATHER_DICT
assert len(response["messages"]) == 4
def test_json_schema(self) -> None:
"""Test response_format as NativeOutput with JSON schema."""
tool_calls = [
[{"args": {}, "id": "1", "name": "get_weather"}],
]
model = FakeToolCallingModel[dict](
tool_calls=tool_calls, structured_response=EXPECTED_WEATHER_DICT
)
agent = create_agent(
model, [get_weather], response_format=NativeOutput(weather_json_schema)
)
response = agent.invoke({"messages": [HumanMessage("What's the weather?")]})
assert response["structured_response"] == EXPECTED_WEATHER_DICT
assert len(response["messages"]) == 4
def test_union_of_types() -> None:
"""Test response_format as NativeOutput with Union (if supported)."""
tool_calls = [
[{"args": {}, "id": "1", "name": "get_weather"}],
[
{
"name": "WeatherBaseModel",
"id": "2",
"args": WEATHER_DATA,
}
],
]
model = FakeToolCallingModel[Union[WeatherBaseModel, LocationResponse]](
tool_calls=tool_calls, structured_response=EXPECTED_WEATHER_PYDANTIC
)
agent = create_agent(
model,
[get_weather, get_location],
response_format=ToolOutput(Union[WeatherBaseModel, LocationResponse]),
)
response = agent.invoke({"messages": [HumanMessage("What's the weather?")]})
assert response["structured_response"] == EXPECTED_WEATHER_PYDANTIC
assert len(response["messages"]) == 5
@pytest.mark.skipif(
skip_openai_integration_tests, reason="OpenAI integration tests are disabled."
)
def test_inference_to_native_output() -> None:
"""Test that native output is inferred when a model supports it."""
model = ChatOpenAI(model="gpt-5")
agent = create_agent(
model,
prompt="You are a helpful weather assistant. Please call the get_weather tool, then use the WeatherReport tool to generate the final response.",
tools=[get_weather],
response_format=WeatherBaseModel,
)
response = agent.invoke({"messages": [HumanMessage("What's the weather?")]})
assert isinstance(response["structured_response"], WeatherBaseModel)
assert response["structured_response"].temperature == 75.0
assert response["structured_response"].condition.lower() == "sunny"
assert len(response["messages"]) == 4
assert [m.type for m in response["messages"]] == [
"human", # "What's the weather?"
"ai", # "What's the weather?"
"tool", # "The weather is sunny and 75°F."
"ai", # structured response
]
@pytest.mark.skipif(
skip_openai_integration_tests, reason="OpenAI integration tests are disabled."
)
def test_inference_to_tool_output() -> None:
"""Test that tool output is inferred when a model supports it."""
model = ChatOpenAI(model="gpt-4")
agent = create_agent(
model,
prompt="You are a helpful weather assistant. Please call the get_weather tool, then use the WeatherReport tool to generate the final response.",
tools=[get_weather],
response_format=ToolOutput(WeatherBaseModel),
)
response = agent.invoke({"messages": [HumanMessage("What's the weather?")]})
assert isinstance(response["structured_response"], WeatherBaseModel)
assert response["structured_response"].temperature == 75.0
assert response["structured_response"].condition.lower() == "sunny"
assert len(response["messages"]) == 5
assert [m.type for m in response["messages"]] == [
"human", # "What's the weather?"
"ai", # "What's the weather?"
"tool", # "The weather is sunny and 75°F."
"ai", # structured response
"tool", # artificial tool message
]
+2 -33
View File
@@ -39,7 +39,7 @@ class TestUsingToolStrategy:
"""Test basic UsingToolStrategy creation."""
strategy = ToolOutput(schema=_TestModel)
assert strategy.schema == _TestModel
assert strategy.tool_message_content == "ok!"
assert strategy.tool_message_content is None
assert len(strategy.schema_specs) == 1
def test_multiple_schemas(self):
@@ -130,40 +130,9 @@ class TestOutputToolBinding:
# 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 data to _TestModel"):
tool_binding.parse(tool_args)
def test_parse_payload_invalid_kind(self):
"""Test parsing with invalid kind."""
from unittest.mock import Mock
mock_tool = Mock()
tool_binding = OutputToolBinding(
schema=_TestModel,
schema_kind="invalid_kind", # type: ignore
tool=mock_tool,
)
with pytest.raises(ValueError, match="Unsupported schema kind: invalid_kind"):
tool_binding.parse({"name": "test", "age": 25})
def test_parse_payload_invalid_pydantic_schema(self):
"""Test parsing with invalid schema for pydantic kind."""
from unittest.mock import Mock
mock_tool = Mock()
# Create tool binding with dict schema but pydantic kind
tool_binding = OutputToolBinding(
schema={"type": "object"}, schema_kind="pydantic", tool=mock_tool
)
with pytest.raises(
ValueError, match="Expected Pydantic model class for 'pydantic' kind"
):
tool_binding.parse({"name": "test", "age": 25})
class TestEdgeCases:
"""Test edge cases and error conditions."""
+152
View File
@@ -0,0 +1,152 @@
from __future__ import annotations
import json
from pathlib import Path
from typing import Any, Dict, List, Optional, Sequence, Type, Union
from unittest.mock import MagicMock
import pytest
from langchain_core.messages import HumanMessage
from langchain_core.tools import tool
from pydantic import BaseModel, create_model
from langgraph.prebuilt import create_agent
from langgraph.prebuilt.responses import ToolOutput
try:
from langchain_openai import ChatOpenAI
except ImportError:
skip_openai_integration_tests = True
else:
skip_openai_integration_tests = False
def _load_spec() -> List[Dict[str, Any]]:
with (Path(__file__).parent / "specifications" / "responses.json").open(
"r", encoding="utf-8"
) as f:
return json.load(f)
TEST_CASES = _load_spec()
AGENT_PROMPT = "You are an HR assistant."
EMPLOYEES = [
{"name": "Sabine", "role": "Developer", "department": "IT"},
{"name": "Henrik", "role": "Product Manager", "department": "IT"},
{"name": "Jessica", "role": "HR", "department": "People"},
]
def _make_tool(fn, *, name: str, description: str):
mock = MagicMock(side_effect=lambda *, name: fn(name=name))
InputModel = create_model(f"{name}_input", name=(str, ...))
@tool(name, description=description, args_schema=InputModel)
def _wrapped(name: str):
return mock(name=name)
return {"tool": _wrapped, "mock": mock}
def _build_tool_output_response_format(
response_format_spec: Sequence[Dict[str, Any]],
) -> ToolOutput:
models: List[Type[BaseModel]] = []
keyset_to_tool_name: Dict[frozenset[str], str] = {}
type_map = {
"string": str,
"number": float,
"integer": int,
"boolean": bool,
"object": dict,
"array": list,
}
for idx, schema in enumerate(response_format_spec):
properties = schema["properties"]
required = set(schema["required"])
type_name = schema.get("title") or f"structured_output_format_{idx + 1}"
fields = {}
for k, prop in properties.items():
py_type = type_map.get(prop.get("type"), Any)
fields[k] = (py_type, ...) if k in required else (Optional[py_type], None)
model = create_model(type_name, **fields)
models.append(model)
keyset_to_tool_name[frozenset(required)] = type_name
union_type = Union[tuple(models)]
return ToolOutput(union_type)
@pytest.mark.skipif(
skip_openai_integration_tests, reason="OpenAI integration tests are disabled."
)
@pytest.mark.xfail(
reason="currently failing due to undefined behavior for multiple structured responses."
)
@pytest.mark.parametrize("case", TEST_CASES, ids=[c["name"] for c in TEST_CASES])
def test_responses_integration_matrix(case: Dict[str, Any]) -> None:
def get_employee_role(*, name: str) -> Optional[str]:
for e in EMPLOYEES:
if e["name"] == name:
return e["role"]
return None
def get_employee_department(*, name: str) -> Optional[str]:
for e in EMPLOYEES:
if e["name"] == name:
return e["department"]
return None
role_tool = _make_tool(
get_employee_role,
name="getEmployeeRole",
description="Get the employee role by name",
)
dept_tool = _make_tool(
get_employee_department,
name="getEmployeeDepartment",
description="Get the employee department by name",
)
response_spec = case["responseFormat"]
if isinstance(response_spec, dict):
response_spec = [response_spec]
tool_output = _build_tool_output_response_format(response_spec)
for assertion in case["assertionsByInvocation"]:
prompt: str = assertion["prompt"]
expected_calls: Dict[str, int] = assertion["toolsWithExpectedCalls"]
expected_structured = assertion.get("expectedStructuredResponse")
expected_last_message = assertion.get("expectedLastMessage")
model = ChatOpenAI(
model="gpt-4o-mini",
temperature=0,
)
agent = create_agent(
model,
tools=[role_tool["tool"], dept_tool["tool"]],
prompt=AGENT_PROMPT,
response_format=tool_output,
)
result = agent.invoke({"messages": [HumanMessage(prompt)]})
# TODO: Count LLM calls. JS handles with mock fetch. Could pass in mock http_client?
# Count tool calls
assert role_tool["mock"].call_count == expected_calls["getEmployeeRole"]
assert dept_tool["mock"].call_count == expected_calls["getEmployeeDepartment"]
# Check last message content
last_message = result["messages"][-1]
assert last_message.content == expected_last_message
# Check structured response
structured_response_json = result["structured_response"].model_dump()
assert structured_response_json == expected_structured
print("Passed test for: ", case["name"])