From f4cdeea6ad5dbf836fb777299b1cd9a0c06c59b1 Mon Sep 17 00:00:00 2001 From: Caspar Broekhuizen Date: Thu, 21 Aug 2025 07:43:50 -0700 Subject: [PATCH] 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 Co-authored-by: Sydney Runkle <54324534+sydney-runkle@users.noreply.github.com> --- .../langgraph/prebuilt/chat_agent_executor.py | 175 ++++-- libs/prebuilt/langgraph/prebuilt/responses.py | 280 +++++++--- libs/prebuilt/tests/model.py | 55 +- .../tests/specifications/responses.json | 46 ++ libs/prebuilt/tests/test_react_agent.py | 6 +- libs/prebuilt/tests/test_response_format.py | 507 ++++++++++++++++++ libs/prebuilt/tests/test_responses.py | 35 +- libs/prebuilt/tests/test_responses_int.py | 152 ++++++ 8 files changed, 1097 insertions(+), 159 deletions(-) create mode 100644 libs/prebuilt/tests/specifications/responses.json create mode 100644 libs/prebuilt/tests/test_response_format.py create mode 100644 libs/prebuilt/tests/test_responses_int.py diff --git a/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py b/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py index e31d99fdd..575c0c6a0 100644 --- a/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py +++ b/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py @@ -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, diff --git a/libs/prebuilt/langgraph/prebuilt/responses.py b/libs/prebuilt/langgraph/prebuilt/responses.py index 9c908ce2e..cd45cc53d 100644 --- a/libs/prebuilt/langgraph/prebuilt/responses.py +++ b/libs/prebuilt/langgraph/prebuilt/responses.py @@ -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]] diff --git a/libs/prebuilt/tests/model.py b/libs/prebuilt/tests/model.py index 2d6f3b773..aab7885b9 100644 --- a/libs/prebuilt/tests/model.py +++ b/libs/prebuilt/tests/model.py @@ -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]], diff --git a/libs/prebuilt/tests/specifications/responses.json b/libs/prebuilt/tests/specifications/responses.json new file mode 100644 index 000000000..7c0cbe4b4 --- /dev/null +++ b/libs/prebuilt/tests/specifications/responses.json @@ -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 + } + ] + } + + ] \ No newline at end of file diff --git a/libs/prebuilt/tests/test_react_agent.py b/libs/prebuilt/tests/test_react_agent.py index 9674e36f4..e2fa81749 100644 --- a/libs/prebuilt/tests/test_react_agent.py +++ b/libs/prebuilt/tests/test_react_agent.py @@ -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", ), diff --git a/libs/prebuilt/tests/test_response_format.py b/libs/prebuilt/tests/test_response_format.py new file mode 100644 index 000000000..26f08e14f --- /dev/null +++ b/libs/prebuilt/tests/test_response_format.py @@ -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 + ] diff --git a/libs/prebuilt/tests/test_responses.py b/libs/prebuilt/tests/test_responses.py index 6fe2f87df..1f191f121 100644 --- a/libs/prebuilt/tests/test_responses.py +++ b/libs/prebuilt/tests/test_responses.py @@ -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.""" diff --git a/libs/prebuilt/tests/test_responses_int.py b/libs/prebuilt/tests/test_responses_int.py new file mode 100644 index 000000000..c4e3363fe --- /dev/null +++ b/libs/prebuilt/tests/test_responses_int.py @@ -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"])