mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-26 19:45:00 +02:00
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:
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user