From 77d98b426b98750a5e5696362e4304cd9ec3f72f Mon Sep 17 00:00:00 2001 From: Caspar Broekhuizen Date: Fri, 22 Aug 2025 10:55:09 -0700 Subject: [PATCH] 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 | 32 +++- .../tests/specifications/responses.json | 43 ++++- .../tests/specifications/return_direct.json | 48 ++++++ libs/prebuilt/tests/test_response_format.py | 72 ++++++++ libs/prebuilt/tests/test_responses_int.py | 152 ----------------- libs/prebuilt/tests/test_responses_spec.py | 160 ++++++++++++++++++ .../prebuilt/tests/test_return_direct_spec.py | 117 +++++++++++++ libs/prebuilt/tests/utils.py | 22 +++ 8 files changed, 484 insertions(+), 162 deletions(-) create mode 100644 libs/prebuilt/tests/specifications/return_direct.json delete mode 100644 libs/prebuilt/tests/test_responses_int.py create mode 100644 libs/prebuilt/tests/test_responses_spec.py create mode 100644 libs/prebuilt/tests/test_return_direct_spec.py create mode 100644 libs/prebuilt/tests/utils.py diff --git a/libs/prebuilt/langgraph/prebuilt/responses.py b/libs/prebuilt/langgraph/prebuilt/responses.py index cd45cc53d..59a0e1177 100644 --- a/libs/prebuilt/langgraph/prebuilt/responses.py +++ b/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 diff --git a/libs/prebuilt/tests/specifications/responses.json b/libs/prebuilt/tests/specifications/responses.json index 7c0cbe4b4..2755ef975 100644 --- a/libs/prebuilt/tests/specifications/responses.json +++ b/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 + } + ] } - ] \ No newline at end of file diff --git a/libs/prebuilt/tests/specifications/return_direct.json b/libs/prebuilt/tests/specifications/return_direct.json new file mode 100644 index 000000000..e4cc5947c --- /dev/null +++ b/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 + } +] \ No newline at end of file diff --git a/libs/prebuilt/tests/test_response_format.py b/libs/prebuilt/tests/test_response_format.py index 26f08e14f..c9d1f8357 100644 --- a/libs/prebuilt/tests/test_response_format.py +++ b/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 diff --git a/libs/prebuilt/tests/test_responses_int.py b/libs/prebuilt/tests/test_responses_int.py deleted file mode 100644 index c4e3363fe..000000000 --- a/libs/prebuilt/tests/test_responses_int.py +++ /dev/null @@ -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"]) diff --git a/libs/prebuilt/tests/test_responses_spec.py b/libs/prebuilt/tests/test_responses_spec.py new file mode 100644 index 000000000..85e146853 --- /dev/null +++ b/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 diff --git a/libs/prebuilt/tests/test_return_direct_spec.py b/libs/prebuilt/tests/test_return_direct_spec.py new file mode 100644 index 000000000..50fb2f6eb --- /dev/null +++ b/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: " 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 diff --git a/libs/prebuilt/tests/utils.py b/libs/prebuilt/tests/utils.py new file mode 100644 index 000000000..957302acf --- /dev/null +++ b/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]