diff --git a/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py b/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py index f54b7b2e9..ac2f27ac0 100644 --- a/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py +++ b/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py @@ -5,10 +5,10 @@ from typing import ( Any, Awaitable, Callable, + Generic, Optional, Sequence, Type, - TypeVar, Union, cast, get_type_hints, @@ -33,7 +33,7 @@ from langchain_core.runnables import ( ) from langchain_core.tools import BaseTool from pydantic import BaseModel -from typing_extensions import Annotated, NotRequired, TypedDict +from typing_extensions import Annotated, NotRequired, TypedDict, TypeVar from langgraph._internal._runnable import RunnableCallable, RunnableLike from langgraph._internal._typing import MISSING @@ -48,7 +48,6 @@ from langgraph.prebuilt._internal._typing import ( from langgraph.prebuilt.responses import ( OutputToolBinding, ResponseFormat, - SchemaSpec, ToolOutput, ) from langgraph.prebuilt.tool_node import ToolNode @@ -58,11 +57,12 @@ from langgraph.types import Checkpointer, Command, Send from langgraph.typing import ContextT from langgraph.warnings import LangGraphDeprecatedSinceV10 -StructuredResponse = Union[dict, BaseModel] -StructuredResponseSchema = Union[dict, type[BaseModel]] - F = TypeVar("F", bound=Callable[..., Any]) +StructuredResponseT = TypeVar( + "StructuredResponseT", bound=Union[dict, BaseModel, None], default=None +) + # We create the AgentState that we will pass around # This simply involves a list of messages @@ -84,16 +84,18 @@ class AgentStatePydantic(BaseModel): remaining_steps: RemainingSteps = 25 -class AgentStateWithStructuredResponse(AgentState): +class AgentStateWithStructuredResponse(AgentState, Generic[StructuredResponseT]): """The state of the agent with a structured response.""" - structured_response: StructuredResponse + structured_response: StructuredResponseT -class AgentStateWithStructuredResponsePydantic(AgentStatePydantic): +class AgentStateWithStructuredResponsePydantic( + AgentStatePydantic, Generic[StructuredResponseT] +): """The state of the agent with a structured response.""" - structured_response: StructuredResponse + structured_response: StructuredResponseT StateSchema = TypeVar("StateSchema", bound=Union[AgentState, AgentStatePydantic]) @@ -185,7 +187,7 @@ def _validate_chat_history( raise ValueError(error_message) -class _AgentBuilder: +class _AgentBuilder(Generic[StructuredResponseT]): """Internal builder class for constructing and agent.""" def __init__( @@ -198,7 +200,7 @@ class _AgentBuilder: tools: Union[Sequence[Union[BaseTool, Callable, dict[str, Any]]], ToolNode], *, prompt: Optional[Prompt] = None, - response_format: Optional[ResponseFormat] = None, + response_format: Optional[ResponseFormat[StructuredResponseT]] = None, pre_model_hook: Optional[RunnableLike] = None, post_model_hook: Optional[RunnableLike] = None, state_schema: Optional[StateSchemaType] = None, @@ -256,15 +258,15 @@ class _AgentBuilder: Future strategies (json_mode, guided) will have separate setup methods. """ - self.structured_output_tools: dict[str, OutputToolBinding] = {} + self.structured_output_tools: dict[ + str, OutputToolBinding[StructuredResponseT] + ] = {} if self.response_format is not None: response_format = self.response_format - # Handle UsingToolStrategy wrapper if isinstance(response_format, ToolOutput): - # Use tools strategy - process each ResponseSchema in the UsingToolStrategy - for response_schema in response_format.schemas: - # Use the factory method to create OutputToolBinding + # check if response_format.schema is a union + for response_schema in response_format.schema_specs: structured_tool_info = OutputToolBinding.from_schema_spec( response_schema ) @@ -275,7 +277,7 @@ class _AgentBuilder: # This shouldn't happen with the new ResponseFormat type, but keeping for safety raise ValueError( f"Unsupported response_format type: {type(response_format)}. " - f"Expected UsingToolStrategy." + f"Expected ToolOutput." ) def _setup_state_schema(self) -> None: @@ -328,12 +330,17 @@ class _AgentBuilder: "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="ok!", + content=tool_message_content, tool_call_id=tool_call["id"], name=tool_call["name"], ), @@ -383,10 +390,8 @@ class _AgentBuilder: if len(all_tools) > 0: # Check if we need to force tool use for structured output tool_choice = None - if ( - self.response_format is not None - and isinstance(self.response_format, ToolOutput) - and self.response_format.tool_choice == "required" + if self.response_format is not None and isinstance( + self.response_format, ToolOutput ): tool_choice = "any" @@ -778,7 +783,7 @@ def create_react_agent( tools: Union[Sequence[Union[BaseTool, Callable, dict[str, Any]]], ToolNode], *, prompt: Optional[Prompt] = None, - response_format: Optional[Union[ToolOutput, StructuredResponseSchema]] = None, + response_format: Optional[Union[ToolOutput, type[StructuredResponseT]]] = None, pre_model_hook: Optional[RunnableLike] = None, post_model_hook: Optional[RunnableLike] = None, state_schema: Optional[StateSchemaType] = None, @@ -985,9 +990,9 @@ def create_react_agent( # 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( - schemas=[SchemaSpec(response_format)], - tool_choice="required", + schema=response_format, ) + elif isinstance(response_format, tuple): if len(response_format) == 2: raise ValueError( diff --git a/libs/prebuilt/langgraph/prebuilt/responses.py b/libs/prebuilt/langgraph/prebuilt/responses.py index 3f6560413..9c908ce2e 100644 --- a/libs/prebuilt/langgraph/prebuilt/responses.py +++ b/libs/prebuilt/langgraph/prebuilt/responses.py @@ -2,33 +2,39 @@ from __future__ import annotations +import sys from dataclasses import dataclass -from typing import Any, Generic, Literal, Optional, Sequence, Type, TypeVar +from typing import Any, Generic, Literal, TypeVar, Union, cast, 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 # For now, we support only Pydantic models as schemas. -SchemaT = TypeVar("SchemaT", bound=type[BaseModel]) -Schema = TypeVar("Schema", bound=BaseModel) +SchemaT = TypeVar("SchemaT") -ToolChoice = Literal["required", "auto"] -"""Required: model must call a tool; auto: model may respond with free text.""" + +if sys.version_info >= (3, 10): + from types import UnionType +else: + UnionType = Union @dataclass(init=False) -class SchemaSpec(Generic[SchemaT]): +class _SchemaSpec(Generic[SchemaT]): """Describes a structured output schema.""" - schema: SchemaT + schema: type[SchemaT] """The schema for the response, can be a dict or a Pydantic model.""" - name: Optional[str] = None - """Name of the schema, used for tool calling. + + name: str | None = None + """Name of the schema, used for tool calling. If not provided, the name will be the model name. """ - description: Optional[str] = None + + description: str | None = None """Custom description of the schema. If not provided, provided will use the model's docstring. @@ -38,10 +44,10 @@ class SchemaSpec(Generic[SchemaT]): def __init__( self, - schema: SchemaT, + schema: type[SchemaT], *, - name: Optional[str] = None, - description: Optional[str] = None, + name: str | None = None, + description: str | None = None, strict: bool = False, ) -> None: """Initialize SchemaSpec with schema and optional parameters.""" @@ -52,31 +58,33 @@ class SchemaSpec(Generic[SchemaT]): @dataclass(init=False) -class ToolOutput: +class ToolOutput(Generic[SchemaT]): """Use a tool calling strategy for model responses.""" - schemas: Sequence[SchemaSpec] - """Schemas for the tool calls.""" - tool_choice: ToolChoice - """Whether to require tool calling or allow the model to choose. - - - "required": The model must use tool calling. - - "auto": The model can choose whether to use tool calling. - - Use `auto` if you want the agent to be able to respond with a non structured - response to ask clarifying questions. - """ + schema: type[SchemaT] + """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.""" def __init__( - self, schemas: Sequence[SchemaSpec], *, tool_choice: ToolChoice = "required" + self, schema: type[SchemaT], tool_message_content: str = "ok!" ) -> None: - """Initialize ToolOutput with schemas and tool choice.""" - self.schemas = schemas - self.tool_choice = tool_choice + """Initialize ToolOutput with schemas and tool message content.""" + 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)] @dataclass -class OutputToolBinding(Generic[Schema]): +class OutputToolBinding(Generic[SchemaT]): """Information for tracking structured output tool metadata. This contains all necessary information to handle structured responses @@ -84,17 +92,17 @@ class OutputToolBinding(Generic[Schema]): and the corresponding tool implementation used by the tools strategy. """ - schema: Type[Schema] + schema: type[SchemaT] """The original schema provided for structured output (Pydantic model or dict schema).""" + schema_kind: Literal["pydantic"] """Classification of the schema type for proper response construction.""" + tool: BaseTool """LangChain tool instance created from the schema for model binding.""" @classmethod - def from_schema_spec( - cls, schema_spec: SchemaSpec[Type[Schema]] - ) -> "OutputToolBinding": + def from_schema_spec(cls, schema_spec: _SchemaSpec[SchemaT]) -> Self: """Create an OutputToolBinding instance from a SchemaSpec. Args: @@ -129,12 +137,12 @@ class OutputToolBinding(Generic[Schema]): tool = tool_creator(schema, **kwargs) return cls( - schema=schema, + schema=cast(type[SchemaT], schema), schema_kind=schema_kind, tool=tool, ) - def parse(self, tool_args: dict[str, Any]) -> Schema: + def parse(self, tool_args: dict[str, Any]) -> SchemaT: """Parse tool arguments according to the schema. Args: @@ -164,4 +172,4 @@ class OutputToolBinding(Generic[Schema]): # TODO: Add support for built-in structured responses (e.g., openai, grok) -ResponseFormat = ToolOutput +ResponseFormat = ToolOutput[SchemaT] diff --git a/libs/prebuilt/tests/model.py b/libs/prebuilt/tests/model.py index 7daf80431..2d6f3b773 100644 --- a/libs/prebuilt/tests/model.py +++ b/libs/prebuilt/tests/model.py @@ -2,6 +2,7 @@ from typing import ( Any, Callable, Dict, + Generic, List, Literal, Optional, @@ -22,12 +23,12 @@ from langchain_core.runnables import Runnable, RunnableLambda from langchain_core.tools import BaseTool from pydantic import BaseModel -from langgraph.prebuilt.chat_agent_executor import StructuredResponse +from langgraph.prebuilt.chat_agent_executor import StructuredResponseT -class FakeToolCallingModel(BaseChatModel): +class FakeToolCallingModel(BaseChatModel, Generic[StructuredResponseT]): tool_calls: Optional[list[list[ToolCall]]] = None - structured_response: Optional[StructuredResponse] = None + structured_response: Optional[StructuredResponseT] = None index: int = 0 tool_style: Literal["openai", "anthropic"] = "openai" @@ -57,7 +58,7 @@ class FakeToolCallingModel(BaseChatModel): def with_structured_output( self, schema: Type[BaseModel] - ) -> Runnable[LanguageModelInput, StructuredResponse]: + ) -> Runnable[LanguageModelInput, StructuredResponseT]: if self.structured_response is None: raise ValueError("Structured response is not set") diff --git a/libs/prebuilt/tests/test_react_agent.py b/libs/prebuilt/tests/test_react_agent.py index d95d43944..b2887d716 100644 --- a/libs/prebuilt/tests/test_react_agent.py +++ b/libs/prebuilt/tests/test_react_agent.py @@ -433,7 +433,7 @@ def test_react_agent_with_structured_response() -> None: return "The weather is sunny and 75°F." expected_structured_response = WeatherResponse(temperature=75) - model = FakeToolCallingModel( + model = FakeToolCallingModel[WeatherResponse]( tool_calls=tool_calls, structured_response=expected_structured_response ) agent = create_react_agent( diff --git a/libs/prebuilt/tests/test_responses.py b/libs/prebuilt/tests/test_responses.py index d6f58cc94..6fe2f87df 100644 --- a/libs/prebuilt/tests/test_responses.py +++ b/libs/prebuilt/tests/test_responses.py @@ -1,17 +1,18 @@ """Unit tests for langgraph.prebuilt.responses module.""" +from typing import Union + import pytest from pydantic import BaseModel from langgraph.prebuilt.responses import ( OutputToolBinding, - ResponseFormat, - SchemaSpec, ToolOutput, + _SchemaSpec, ) -class TestModel(BaseModel): +class _TestModel(BaseModel): """A test model for structured output.""" name: str @@ -31,54 +32,29 @@ class EmptyDocModel(BaseModel): data: str -class TestSchemaSpec: - """Test SchemaSpec dataclass.""" - - def test_basic_creation(self): - """Test basic SchemaSpec creation.""" - schema = SchemaSpec(schema=TestModel) - assert schema.schema == TestModel - assert schema.name is None - assert schema.description is None - assert schema.strict is False - - def test_creation_with_all_fields(self): - """Test SchemaSpec creation with all fields.""" - schema = SchemaSpec( - schema=TestModel, - name="custom_test_model", - description="A custom description", - strict=True, - ) - assert schema.schema == TestModel - assert schema.name == "custom_test_model" - assert schema.description == "A custom description" - assert schema.strict is True - - class TestUsingToolStrategy: """Test UsingToolStrategy dataclass.""" def test_basic_creation(self): """Test basic UsingToolStrategy creation.""" - schema = SchemaSpec(schema=TestModel) - strategy = ToolOutput(schemas=[schema]) - assert len(strategy.schemas) == 1 - assert strategy.schemas[0] == schema - assert strategy.tool_choice == "required" # default - - def test_creation_with_auto_tool_choice(self): - """Test UsingToolStrategy creation with auto tool choice.""" - schema = SchemaSpec(schema=TestModel) - strategy = ToolOutput(schemas=[schema], tool_choice="auto") - assert strategy.tool_choice == "auto" + strategy = ToolOutput(schema=_TestModel) + assert strategy.schema == _TestModel + assert strategy.tool_message_content == "ok!" + assert len(strategy.schema_specs) == 1 def test_multiple_schemas(self): """Test UsingToolStrategy with multiple schemas.""" - schema1 = SchemaSpec(schema=TestModel) - schema2 = SchemaSpec(schema=CustomModel) - strategy = ToolOutput(schemas=[schema1, schema2]) - assert len(strategy.schemas) == 2 + strategy = ToolOutput(schema=Union[_TestModel, CustomModel]) + assert len(strategy.schema_specs) == 2 + assert strategy.schema_specs[0].schema == _TestModel + assert strategy.schema_specs[1].schema == CustomModel + + def test_schema_with_tool_message_content(self): + """Test UsingToolStrategy with tool message content.""" + strategy = ToolOutput(schema=_TestModel, tool_message_content="custom message") + assert strategy.schema == _TestModel + assert strategy.tool_message_content == "custom message" + assert len(strategy.schema_specs) == 1 class TestOutputToolBinding: @@ -86,24 +62,24 @@ class TestOutputToolBinding: def test_from_schema_spec_basic(self): """Test basic OutputToolBinding creation from SchemaSpec.""" - schema_spec = SchemaSpec(schema=TestModel) + schema_spec = _SchemaSpec(schema=_TestModel) tool_binding = OutputToolBinding.from_schema_spec(schema_spec) - assert tool_binding.schema == TestModel + assert tool_binding.schema == _TestModel assert tool_binding.schema_kind == "pydantic" assert tool_binding.tool is not None - assert tool_binding.tool.name == "TestModel" + assert tool_binding.tool.name == "_TestModel" def test_from_schema_spec_with_custom_name(self): """Test OutputToolBinding creation with custom name.""" - schema_spec = SchemaSpec(schema=TestModel, name="custom_tool_name") + schema_spec = _SchemaSpec(schema=_TestModel, name="custom_tool_name") tool_binding = OutputToolBinding.from_schema_spec(schema_spec) assert tool_binding.tool.name == "custom_tool_name" def test_from_schema_spec_with_custom_description(self): """Test OutputToolBinding creation with custom description.""" - schema_spec = SchemaSpec( - schema=TestModel, description="Custom tool description" + schema_spec = _SchemaSpec( + schema=_TestModel, description="Custom tool description" ) tool_binding = OutputToolBinding.from_schema_spec(schema_spec) @@ -111,7 +87,7 @@ class TestOutputToolBinding: def test_from_schema_spec_with_model_docstring(self): """Test OutputToolBinding creation using model docstring as description.""" - schema_spec = SchemaSpec(schema=CustomModel) + schema_spec = _SchemaSpec(schema=CustomModel) tool_binding = OutputToolBinding.from_schema_spec(schema_spec) assert tool_binding.tool.description == "Custom model with a custom docstring." @@ -127,7 +103,7 @@ class TestOutputToolBinding: # This should have the same docstring as BaseModel pass - schema_spec = SchemaSpec(schema=DefaultDocModel) + schema_spec = _SchemaSpec(schema=DefaultDocModel) tool_binding = OutputToolBinding.from_schema_spec(schema_spec) # Should use empty description when model has default BaseModel docstring @@ -135,26 +111,26 @@ class TestOutputToolBinding: def test_parse_payload_pydantic_success(self): """Test successful parsing for Pydantic model.""" - schema_spec = SchemaSpec(schema=TestModel) + schema_spec = _SchemaSpec(schema=_TestModel) tool_binding = OutputToolBinding.from_schema_spec(schema_spec) tool_args = {"name": "John", "age": 30} result = tool_binding.parse(tool_args) - assert isinstance(result, TestModel) + assert isinstance(result, _TestModel) assert result.name == "John" assert result.age == 30 assert result.email == "default@example.com" # default value def test_parse_payload_pydantic_validation_error(self): """Test parsing failure for invalid Pydantic data.""" - schema_spec = SchemaSpec(schema=TestModel) + schema_spec = _SchemaSpec(schema=_TestModel) tool_binding = OutputToolBinding.from_schema_spec(schema_spec) # 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 tool args to _TestModel"): tool_binding.parse(tool_args) def test_parse_payload_invalid_kind(self): @@ -164,7 +140,7 @@ class TestOutputToolBinding: mock_tool = Mock() tool_binding = OutputToolBinding( - schema=TestModel, + schema=_TestModel, schema_kind="invalid_kind", # type: ignore tool=mock_tool, ) @@ -189,36 +165,20 @@ class TestOutputToolBinding: tool_binding.parse({"name": "test", "age": 25}) -class TestResponseFormat: - """Test ResponseFormat type alias.""" - - def test_response_format_is_using_tool_strategy(self): - """Test that ResponseFormat is aliased to UsingToolStrategy.""" - assert ResponseFormat is ToolOutput - - def test_can_create_response_format(self): - """Test that we can create ResponseFormat instances.""" - schema = SchemaSpec(schema=TestModel) - response_format = ResponseFormat(schemas=[schema]) - - assert isinstance(response_format, ToolOutput) - assert len(response_format.schemas) == 1 - - class TestEdgeCases: """Test edge cases and error conditions.""" def test_empty_schemas_list(self) -> None: """Test UsingToolStrategy with empty schemas list.""" - strategy = ToolOutput([SchemaSpec(EmptyDocModel)]) - assert len(strategy.schemas) == 1 + strategy = ToolOutput(EmptyDocModel) + assert len(strategy.schema_specs) == 1 @pytest.mark.skip( reason="Need to fix bug in langchain-core for inheritance of doc-strings." ) def test_base_model_doc_constant(self) -> None: """Test that BASE_MODEL_DOC constant is set correctly.""" - binding = OutputToolBinding.from_schema_spec(SchemaSpec(EmptyDocModel)) + binding = OutputToolBinding.from_schema_spec(_SchemaSpec(EmptyDocModel)) assert binding.tool.name == "EmptyDocModel" assert ( binding.tool.description[:5] == ""