fix(langgraph): handle multiple annotations w/ BaseChannel detection (#6210)

This PR ensures that even if a type has multiple annotations, we can
still detect the `BaseChannel` subclasses attached.

```py
class State(TypedDict):
    # recognized as EphemeralValue(int)
    foo: Annotated[int, EphemeralValue]

    # now recognized as EphemeralValue(int)
    bar: Annotated[int, EphemeralValue, OtherMetadata]

    # now recognized as EphemeralValue(int)
    baz: Annotated[int, SomeMetadata, EphemeralValue, OtherMetadata]
```
This commit is contained in:
Sydney Runkle
2025-09-26 17:22:24 -04:00
committed by GitHub
parent 20ddb2b8b4
commit 36179ab1d2
2 changed files with 46 additions and 6 deletions
+8 -4
View File
@@ -1369,10 +1369,14 @@ def _get_channel(
def _is_field_channel(typ: type[Any]) -> BaseChannel | None:
if hasattr(typ, "__metadata__"):
meta = typ.__metadata__
if len(meta) >= 1 and isinstance(meta[-1], BaseChannel):
return meta[-1]
elif len(meta) >= 1 and isclass(meta[-1]) and issubclass(meta[-1], BaseChannel):
return meta[-1](typ.__origin__ if hasattr(typ, "__origin__") else typ)
# Search through all annotated medata to find channel annotations
for item in meta:
if isinstance(item, BaseChannel):
return item
elif isclass(item) and issubclass(item, BaseChannel):
# ex, Annotated[int, EphemeralValue, SomeOtherAnnotation]
# would return EphemeralValue(int)
return item(typ.__origin__ if hasattr(typ, "__origin__") else typ)
return None
+38 -2
View File
@@ -2,7 +2,7 @@ import inspect
import operator
import warnings
from dataclasses import dataclass, field
from typing import Annotated, Any, Optional
from typing import Annotated, Any, Optional, Union
from typing import Annotated as Annotated2
import pytest
@@ -11,7 +11,13 @@ from pydantic import BaseModel
from typing_extensions import NotRequired, Required, TypedDict
from langgraph.channels.binop import BinaryOperatorAggregate
from langgraph.graph.state import StateGraph, _get_node_name, _warn_invalid_state_schema
from langgraph.channels.ephemeral_value import EphemeralValue
from langgraph.graph.state import (
StateGraph,
_get_node_name,
_is_field_channel,
_warn_invalid_state_schema,
)
class State(BaseModel):
@@ -335,3 +341,33 @@ def test_private_input_schema_conditional_edge():
builder.add_edge("__start__", "node_1")
graph = builder.compile()
assert graph.invoke({"foo": 0}) == {"foo": 2, "bar": "meow"}
def test_is_field_channel() -> None:
"""Test channel detection across all scenarios."""
# Basic detection
result = _is_field_channel(Annotated[int, EphemeralValue])
assert isinstance(result, EphemeralValue) and result.typ is int
# Main fix: handles extraneous annotations
result = _is_field_channel(Annotated[str, "metadata", EphemeralValue, "more"])
assert isinstance(result, EphemeralValue) and result.typ is str
# Complex types work
union_type = Union[int, str]
result = _is_field_channel(Annotated[union_type, EphemeralValue])
assert isinstance(result, EphemeralValue) and result.typ is union_type
# Pre-instantiated channels
instantiated = EphemeralValue(int)
result = _is_field_channel(Annotated[int, instantiated])
assert result is instantiated
# Pre-instantiated channels with multiple annotations
instantiated = EphemeralValue(int)
result = _is_field_channel(Annotated[int, "metadata", instantiated, "more"])
assert result is instantiated
# No channel cases
assert _is_field_channel(int) is None
assert _is_field_channel(Annotated[int, "just_metadata"]) is None