mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-08 10:47:52 +02:00
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:
co-authored by
Sydney Runkle
Sydney Runkle
parent
1cd1373788
commit
f4cdeea6ad
@@ -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,
|
||||
|
||||
@@ -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]]
|
||||
|
||||
@@ -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
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
]
|
||||
@@ -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",
|
||||
),
|
||||
|
||||
@@ -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
|
||||
]
|
||||
@@ -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."""
|
||||
|
||||
@@ -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"])
|
||||
Reference in New Issue
Block a user