test(prebuilt): standard integration tests for create_agent (#5988)

This PR improves standard integration tests in prebuilt to ensure
logical equivalence between the Python and JavaScript implementations of
`create_agent`.
* Cleans up `test_responses_int` test harness
* Adds new test utils to dynamically load JSON test specs
* Adds `test_return_direct_int` test harness to validate model behavior
when `return_direct` tool property is set
* Adds support for when the user instantiates `ToolOutput` with multiple
JSON schemas unified by the `oneOf` keyword

Failing tests:
* `test_inference_to_native_output`: there is some odd behavior where
the model makes a second call to `get_weather` despite having just
received the tool message, so there are 6 messages instead of the 4
expected.
* `test_responses_integration_matrix[asking for information that does
not fit into the response format]`: `XFAIL`, currently failing due to
undefined behavior when the model cannot conform to any of the
structured response formats.

TODO in future PRs:
* Add exception handling to pass `test_responses_integration_matrix`.

libs/prebuilt/langgraph/prebuilt/responses.py

@@ -3,8 +3,9 @@
 from __future__ import annotations
 
 import sys
+import uuid
 from dataclasses import dataclass, is_dataclass
-from typing import Any, Generic, Literal, TypeVar, Union, get_args, get_origin
+from typing import Any, Generic, Iterable, Literal, TypeVar, Union, get_args, get_origin
 
 from langchain_core.messages import AIMessage
 from langchain_core.tools import BaseTool, StructuredTool
@@ -59,7 +60,7 @@ class _SchemaSpec(Generic[SchemaT]):
     name: str
     """Name of the schema, used for tool calling.
     
-    If not provided, the name will be the model name or "structured_output" if it's a JSON schema.
+    If not provided, the name will be the model name or "response_format" if it's a JSON schema.
     """
 
     description: str
@@ -88,10 +89,11 @@ class _SchemaSpec(Generic[SchemaT]):
         """Initialize SchemaSpec with schema and optional parameters."""
         self.schema = schema
 
+        # Schema names must be unique so we use a shortened UUID suffix
         self.name = name or (
-            schema.get("title", "structured_output")
+            schema.get("title", f"response_format_{str(uuid.uuid4())[:4]}")
             if isinstance(schema, dict)
-            else getattr(schema, "__name__", "structured_output")
+            else getattr(schema, "__name__", f"response_format_{str(uuid.uuid4())[:4]}")
         )
 
         self.description = description or (
@@ -143,10 +145,22 @@ class ToolOutput(Generic[SchemaT]):
         self.schema = schema
         self.tool_message_content = tool_message_content
 
-        if get_origin(schema) in (UnionType, Union):
-            self.schema_specs = [_SchemaSpec(s) for s in get_args(schema)]
-        else:
-            self.schema_specs = [_SchemaSpec(schema)]
+        def _iter_variants(schema: Any) -> Iterable[Any]:
+            """Yield leaf variants from Union and JSON Schema oneOf."""
+
+            if get_origin(schema) in (UnionType, Union):
+                for arg in get_args(schema):
+                    yield from _iter_variants(arg)
+                return
+
+            if isinstance(schema, dict) and "oneOf" in schema:
+                for sub in schema.get("oneOf", []):
+                    yield from _iter_variants(sub)
+                return
+
+            yield schema
+
+        self.schema_specs = [_SchemaSpec(s) for s in _iter_variants(schema)]
 
 
 @dataclass(init=False)
@@ -282,7 +296,7 @@ class NativeOutputBinding(Generic[SchemaT]):
         try:
             data = json.loads(raw_text)
         except Exception as e:
-            schema_name = getattr(self.schema, "__name__", "structured_output")
+            schema_name = getattr(self.schema, "__name__", "response_format")
             raise ValueError(
                 f"Native structured output expected valid JSON for {schema_name}, but parsing failed: {e}."
             ) from e

libs/prebuilt/tests/specifications/responses.json

@@ -3,6 +3,7 @@
       "name": "updated structured response",
       "responseFormat": [
         {
+          "title": "role_schema_structured_output",
           "type": "object",
           "properties": {
             "name": { "type": "string" },
@@ -11,6 +12,7 @@
           "required": ["name", "role"]
         },
         {
+          "title": "department_schema_structured_output",
           "type": "object",
           "properties": {
             "name": { "type": "string" },
@@ -41,6 +43,45 @@
           "llmRequestCount": 4
         }
       ]
+    },
+    {
+      "name": "asking for information that does not fit into the response format",
+      "responseFormat": [
+        {
+          "schema": {
+            "type": "object",
+            "properties": {
+              "name": { "type": "string" },
+              "role": { "type": "string" }
+            },
+            "required": ["name", "role"]
+          }
+        },
+        {
+          "schema": {
+            "type": "object",
+            "properties": {
+              "name": { "type": "string" },
+              "department": { "type": "string" }
+            },
+            "required": ["name", "department"]
+          }
+        }
+      ],
+      "assertionsByInvocation": [
+        {
+          "prompt": "How much does Saskia earn?",
+          "toolsWithExpectedCalls": {
+            "getEmployeeRole": 1,
+            "getEmployeeDepartment": 0
+          },
+          "expectedLastMessage": "Returning structured response: {'name': 'Saskia', 'role': 'Software Engineer'}",
+          "expectedStructuredResponse": {
+            "name": "Saskia",
+            "role": "Software Engineer"
+          },
+          "llmRequestCount": 2
+        }
+      ]
     }
-    
   ]

libs/prebuilt/tests/specifications/return_direct.json

@@ -0,0 +1,48 @@
+[
+  {
+    "name": "Scenario: NO return_direct, NO response_format",
+    "returnDirect": false,
+    "responseFormat": null,
+    "expectedToolCalls": 10,
+    "expectedLastMessage": "Attempts: 10",
+    "expectedStructuredResponse": null
+  },
+  {
+    "name": "Scenario: NO return_direct, YES response_format",
+    "returnDirect": false,
+    "responseFormat": {
+      "type": "object",
+      "properties": {
+        "attempts": { "type": "number" },
+        "succeeded": { "type": "boolean" }
+      },
+      "required": ["attempts", "succeeded"]
+    },
+    "expectedToolCalls": 10,
+    "expectedLastMessage": "Returning structured response: {'attempts': 10, 'succeeded': True}",
+    "expectedStructuredResponse": { "attempts": 10, "succeeded": true }
+  },
+  {
+    "name": "Scenario: YES return_direct, NO response_format",
+    "returnDirect": true,
+    "responseFormat": null,
+    "expectedToolCalls": 1,
+    "expectedLastMessage": "{\"status\": \"pending\", \"attempts\": 1}",
+    "expectedStructuredResponse": null
+  },
+  {
+    "name": "Scenario: YES return_direct, YES response_format",
+    "returnDirect": true,
+    "responseFormat": {
+      "type": "object",
+      "properties": {
+        "attempts": { "type": "number" },
+        "succeeded": { "type": "boolean" }
+      },
+      "required": ["attempts", "succeeded"]
+    },
+    "expectedToolCalls": 1,
+    "expectedLastMessage": "{\"status\": \"pending\", \"attempts\": 1}",
+    "expectedStructuredResponse": null
+  }
+]

libs/prebuilt/tests/test_response_format.py

@@ -59,6 +59,22 @@ class LocationResponse(BaseModel):
     country: str = Field(description="The country name")
 
 
+class LocationTypedDict(TypedDict):
+    city: str
+    country: str
+
+
+location_json_schema = {
+    "type": "object",
+    "properties": {
+        "city": {"type": "string", "description": "The city name"},
+        "country": {"type": "string", "description": "The country name"},
+    },
+    "title": "location_schema",
+    "required": ["city", "country"],
+}
+
+
 def get_weather() -> str:
     """Get the weather."""
 
@@ -80,6 +96,7 @@ 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)
+EXPECTED_LOCATION_DICT: LocationTypedDict = {"city": "New York", "country": "USA"}
 
 
 class TestResponseFormatAsModel:
@@ -261,6 +278,61 @@ class TestResponseFormatAsToolOutput:
         assert response["structured_response"] == EXPECTED_WEATHER_DICT
         assert len(response["messages"]) == 5
 
+    def test_union_of_json_schemas(self) -> None:
+        """Test response_format as ToolOutput with union of JSON schemas."""
+        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, get_location],
+            response_format=ToolOutput(
+                {"oneOf": [weather_json_schema, location_json_schema]}
+            ),
+        )
+        response = agent.invoke({"messages": [HumanMessage("What's the weather?")]})
+
+        assert response["structured_response"] == EXPECTED_WEATHER_DICT
+        assert len(response["messages"]) == 5
+
+        # Test with LocationResponse
+        tool_calls_location = [
+            [{"args": {}, "id": "1", "name": "get_location"}],
+            [
+                {
+                    "name": "location_schema",
+                    "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(
+                {"oneOf": [weather_json_schema, location_json_schema]}
+            ),
+        )
+        response_location = agent_location.invoke(
+            {"messages": [HumanMessage("Where am I?")]}
+        )
+
+        assert response_location["structured_response"] == EXPECTED_LOCATION_DICT
+        assert len(response_location["messages"]) == 5
+
     def test_union_of_types(self) -> None:
         """Test response_format as ToolOutput with Union of various types."""
         # Test with WeatherBaseModel

libs/prebuilt/tests/test_responses_int.py

@@ -1,152 +0,0 @@
-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"])

libs/prebuilt/tests/test_responses_spec.py

@@ -0,0 +1,160 @@
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional, Union
+from unittest.mock import MagicMock
+
+import httpx
+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
+from tests.utils import BaseSchema, load_spec
+
+try:
+    from langchain_openai import ChatOpenAI
+except ImportError:
+    skip_openai_integration_tests = True
+else:
+    skip_openai_integration_tests = False
+
+AGENT_PROMPT = "You are an HR assistant."
+
+
+class ToolCalls(BaseSchema):
+    get_employee_role: int
+    get_employee_department: int
+
+
+class AssertionByInvocation(BaseSchema):
+    prompt: str
+    tools_with_expected_calls: ToolCalls
+    expected_last_message: str
+    expected_structured_response: Optional[Dict[str, Any]]
+    llm_request_count: int
+
+
+class TestCase(BaseSchema):
+    name: str
+    response_format: Union[Dict[str, Any], List[Dict[str, Any]]]
+    assertions_by_invocation: List[AssertionByInvocation]
+
+
+class Employee(BaseModel):
+    name: str
+    role: str
+    department: str
+
+
+EMPLOYEES: list[Employee] = [
+    Employee(name="Sabine", role="Developer", department="IT"),
+    Employee(name="Henrik", role="Product Manager", department="IT"),
+    Employee(name="Jessica", role="HR", department="People"),
+]
+
+TEST_CASES = load_spec("responses", as_model=TestCase)
+
+
+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}
+
+
+@pytest.mark.skipif(
+    skip_openai_integration_tests, reason="OpenAI integration tests are disabled."
+)
+@pytest.mark.parametrize("case", TEST_CASES, ids=[c.name for c in TEST_CASES])
+def test_responses_integration_matrix(case: TestCase) -> None:
+    if case.name == "asking for information that does not fit into the response format":
+        pytest.xfail(
+            "currently failing due to undefined behavior when model cannot conform to any of the structured response formats."
+        )
+
+    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="get_employee_role",
+        description="Get the employee role by name",
+    )
+    dept_tool = _make_tool(
+        get_employee_department,
+        name="get_employee_department",
+        description="Get the employee department by name",
+    )
+
+    response_format_spec = case.response_format
+    if isinstance(response_format_spec, dict):
+        response_format_spec = [response_format_spec]
+    # Unwrap nested schema objects
+    response_format_spec = [item.get("schema", item) for item in response_format_spec]
+    if len(response_format_spec) == 1:
+        tool_output = ToolOutput(response_format_spec[0])
+    else:
+        tool_output = ToolOutput({"oneOf": response_format_spec})
+
+    llm_request_count = 0
+
+    for assertion in case.assertions_by_invocation:
+
+        def on_request(request: httpx.Request) -> None:
+            nonlocal llm_request_count
+            llm_request_count += 1
+
+        http_client = httpx.Client(
+            event_hooks={"request": [on_request]},
+        )
+
+        model = ChatOpenAI(
+            model="gpt-4o",
+            temperature=0,
+            http_client=http_client,
+        )
+
+        agent = create_agent(
+            model,
+            tools=[role_tool["tool"], dept_tool["tool"]],
+            prompt=AGENT_PROMPT,
+            response_format=tool_output,
+        )
+
+        result = agent.invoke({"messages": [HumanMessage(assertion.prompt)]})
+
+        # Count tool calls
+        assert (
+            role_tool["mock"].call_count
+            == assertion.tools_with_expected_calls.get_employee_role
+        )
+        assert (
+            dept_tool["mock"].call_count
+            == assertion.tools_with_expected_calls.get_employee_department
+        )
+
+        # Count LLM calls
+        assert llm_request_count == assertion.llm_request_count
+
+        # Check last message content
+        last_message = result["messages"][-1]
+        assert last_message.content == assertion.expected_last_message
+
+        # Check structured response
+        structured_response_json = result["structured_response"]
+        assert structured_response_json == assertion.expected_structured_response

libs/prebuilt/tests/test_return_direct_spec.py

@@ -0,0 +1,117 @@
+from __future__ import annotations
+
+from typing import Any, Dict, Optional
+from unittest.mock import MagicMock
+
+import pytest
+from langchain_core.messages import HumanMessage
+from langchain_core.tools import tool
+
+from langgraph.prebuilt import create_agent
+from langgraph.prebuilt.responses import ToolOutput
+from tests.utils import BaseSchema, load_spec
+
+try:
+    from langchain_openai import ChatOpenAI
+except ImportError:
+    skip_openai_integration_tests = True
+else:
+    skip_openai_integration_tests = False
+
+AGENT_PROMPT = """
+You are a strict polling bot.
+
+- Only use the "poll_job" tool until it returns { status: "succeeded" }.
+- If status is "pending", call the tool again. Do not produce a final answer.
+- When it is "succeeded", return exactly: "Attempts: <number>" with no extra text.
+"""
+
+
+class TestCase(BaseSchema):
+    name: str
+    return_direct: bool
+    response_format: Optional[Dict[str, Any]]
+    expected_tool_calls: int
+    expected_last_message: str
+    expected_structured_response: Optional[Dict[str, Any]]
+
+
+TEST_CASES = load_spec("return_direct", as_model=TestCase)
+
+
+def _make_tool(return_direct: bool):
+    attempts = 0
+
+    def _side_effect():
+        nonlocal attempts
+        attempts += 1
+        return {
+            "status": "succeeded" if attempts >= 10 else "pending",
+            "attempts": attempts,
+        }
+
+    mock = MagicMock(side_effect=_side_effect)
+
+    @tool(
+        "pollJob",
+        description=(
+            "Check the status of a long-running job. "
+            "Returns { status: 'pending' | 'succeeded', attempts: number }."
+        ),
+        return_direct=return_direct,
+    )
+    def _wrapped():
+        return mock()
+
+    return {"tool": _wrapped, "mock": mock}
+
+
+@pytest.mark.skipif(
+    skip_openai_integration_tests, reason="OpenAI integration tests are disabled."
+)
+@pytest.mark.parametrize("case", TEST_CASES, ids=[c.name for c in TEST_CASES])
+def test_return_direct_integration_matrix(case: TestCase) -> None:
+    poll_tool = _make_tool(case.return_direct)
+
+    model = ChatOpenAI(
+        model="gpt-4o",
+        temperature=0,
+    )
+
+    if case.response_format:
+        agent = create_agent(
+            model,
+            tools=[poll_tool["tool"]],
+            prompt=AGENT_PROMPT,
+            response_format=ToolOutput(case.response_format),
+        )
+    else:
+        agent = create_agent(
+            model,
+            tools=[poll_tool["tool"]],
+            prompt=AGENT_PROMPT,
+        )
+
+    result = agent.invoke(
+        {
+            "messages": [
+                HumanMessage(
+                    "Poll the job until it's done and tell me how many attempts it took."
+                )
+            ]
+        }
+    )
+
+    # Count tool calls
+    assert poll_tool["mock"].call_count == case.expected_tool_calls
+
+    # Check last message content
+    last_message = result["messages"][-1]
+    assert last_message.content == case.expected_last_message
+
+    # Check structured response
+    if case.expected_structured_response is not None:
+        structured_response_json = result["structured_response"]
+        assert structured_response_json == case.expected_structured_response
+    else:
+        assert "structured_response" not in result

libs/prebuilt/tests/utils.py

@@ -0,0 +1,22 @@
+import json
+from pathlib import Path
+from typing import Type
+
+from pydantic import BaseModel, ConfigDict
+from pydantic.alias_generators import to_camel
+
+
+class BaseSchema(BaseModel):
+    model_config = ConfigDict(
+        alias_generator=to_camel,
+        populate_by_name=True,
+        from_attributes=True,
+    )
+
+
+def load_spec(spec_name: str, as_model: Type[BaseModel]) -> list[BaseModel]:
+    with (Path(__file__).parent / "specifications" / f"{spec_name}.json").open(
+        "r", encoding="utf-8"
+    ) as f:
+        data = json.load(f)
+        return [as_model(**item) for item in data]