mirror of
https://github.com/hwchase17/langchain.git
synced 2026-07-12 11:21:37 +00:00
fix(langchain): ordered schema resolution — list replaces set so state_schema wins (#37223)
## Problem
`_resolve_schemas` collected middleware state schemas into a
`set[type]`, making field-conflict resolution non-deterministic. When
two schemas declare the same field with different annotations (e.g.
`DeltaChannel` vs `BinaryOperatorAggregate` on `messages`), the winner
depended on set iteration order — effectively random across runs.
This forced callers to work around it with post-compilation patches
like:
```python
if not isinstance(agent.channels.get("messages"), DeltaChannel):
agent.channels["messages"] = DeltaChannel(...)
```
## Fix
Change `_resolve_schemas` to accept `list[type]` and build it as:
```python
state_schemas = [*(m.state_schema for m in middleware), base_state]
```
`base_state` (the caller's explicit `state_schema`, or `AgentState` by
default) goes **last** and wins any conflict. Middleware schemas precede
it in registration order.
This means callers can now pass `state_schema=MyCustomState` and be
guaranteed it overrides any conflicting middleware annotation — no
post-compilation patching needed.
## Test plan
- All 9 existing `test_state_schema.py` tests pass
- Updated the two direct `_resolve_schemas({...})` call sites in tests
to use lists
This commit is contained in:
@@ -407,8 +407,13 @@ def _get_schema_type_hints(schema: type) -> dict[str, Any]:
|
||||
return get_type_hints(schema, include_extras=True)
|
||||
|
||||
|
||||
def _resolve_schemas(schemas: set[type]) -> tuple[type, type, type]:
|
||||
"""Resolve state, input, and output schemas for the given schemas."""
|
||||
def _resolve_schemas(schemas: list[type]) -> tuple[type, type, type]:
|
||||
"""Resolve state, input, and output schemas for the given schemas.
|
||||
|
||||
Schemas are merged in list order; later entries override earlier ones when the
|
||||
same field is declared by multiple schemas. Duplicates are harmless — a type
|
||||
that appears more than once is processed at its last position.
|
||||
"""
|
||||
schema_hints = {schema: _get_schema_type_hints(schema) for schema in schemas}
|
||||
return (
|
||||
_resolve_schema(schema_hints, "StateSchema", None),
|
||||
@@ -1030,10 +1035,13 @@ def create_agent(
|
||||
]
|
||||
awrap_model_call_handler = _chain_async_model_call_handlers(async_handlers)
|
||||
|
||||
state_schemas: set[type] = {m.state_schema for m in middleware}
|
||||
# Use provided state_schema if available, otherwise use base AgentState
|
||||
base_state = state_schema if state_schema is not None else AgentState
|
||||
state_schemas.add(base_state)
|
||||
# Build an ordered list: middleware schemas first (in registration order),
|
||||
# base_state last so it wins any field conflict. This lets the caller's
|
||||
# explicit state_schema override middleware annotations — e.g. passing
|
||||
# a DeltaChannel-annotated schema wins over BinaryOperatorAggregate from
|
||||
# AgentState without requiring a post-compilation patch.
|
||||
state_schemas: list[type] = [*(m.state_schema for m in middleware), base_state]
|
||||
|
||||
resolved_state_schema, input_schema, output_schema = _resolve_schemas(state_schemas)
|
||||
|
||||
|
||||
@@ -6,7 +6,7 @@ AgentState without needing to create custom middleware.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Annotated, Any
|
||||
from typing import TYPE_CHECKING, Annotated, Any, get_args, get_type_hints
|
||||
|
||||
from langchain_core.messages import HumanMessage
|
||||
from langchain_core.tools import tool
|
||||
@@ -256,6 +256,27 @@ def test_state_schema_with_private_state_field() -> None:
|
||||
assert len(result["messages"]) == 4 # Human, AI (tool call), Tool result, AI (final)
|
||||
|
||||
|
||||
def test_last_schema_wins_for_conflicting_field() -> None:
|
||||
"""The last schema in the list wins when multiple schemas declare the same field.
|
||||
|
||||
In practice this means state_schema (passed last) overrides middleware annotations,
|
||||
and a later middleware overrides an earlier one.
|
||||
"""
|
||||
|
||||
class FirstState(AgentState[Any]):
|
||||
shared_field: Annotated[str, "first"]
|
||||
|
||||
class MiddleState(AgentState[Any]):
|
||||
shared_field: Annotated[str, "middle"]
|
||||
|
||||
class LastState(AgentState[Any]):
|
||||
shared_field: Annotated[str, "last"]
|
||||
|
||||
resolved, _, _ = factory._resolve_schemas([FirstState, MiddleState, LastState])
|
||||
hints = get_type_hints(resolved, include_extras=True)
|
||||
assert get_args(hints["shared_field"])[1] == "last"
|
||||
|
||||
|
||||
def test_get_schema_type_hints_cache_hits_for_reused_schema() -> None:
|
||||
"""Test repeated schema resolution reuses cached type hints for the same schema."""
|
||||
|
||||
@@ -266,9 +287,9 @@ def test_get_schema_type_hints_cache_hits_for_reused_schema() -> None:
|
||||
|
||||
factory._get_schema_type_hints.cache_clear()
|
||||
|
||||
factory._resolve_schemas({CachedState})
|
||||
factory._resolve_schemas([CachedState])
|
||||
first_info = factory._get_schema_type_hints.cache_info()
|
||||
factory._resolve_schemas({CachedState})
|
||||
factory._resolve_schemas([CachedState])
|
||||
second_info = factory._get_schema_type_hints.cache_info()
|
||||
|
||||
assert first_info.misses == 1
|
||||
@@ -293,9 +314,9 @@ def test_get_schema_type_hints_cache_accepts_distinct_local_schema_types() -> No
|
||||
schema_a = make_state_schema("LocalStateA")
|
||||
schema_b = make_state_schema("LocalStateB")
|
||||
|
||||
factory._resolve_schemas({schema_a, schema_b})
|
||||
factory._resolve_schemas([schema_a, schema_b])
|
||||
first_info = factory._get_schema_type_hints.cache_info()
|
||||
factory._resolve_schemas({schema_a, schema_b})
|
||||
factory._resolve_schemas([schema_a, schema_b])
|
||||
second_info = factory._get_schema_type_hints.cache_info()
|
||||
|
||||
assert first_info.misses == 2
|
||||
|
||||
Reference in New Issue
Block a user