Coverage for core/langgraph/_channels.py: 100%

1""" 

2_channels.py — Core-owned channel-context storage primitives. 

3 

4Channel context is Core-owned: committed channel values live as flat keys inside 

5the reserved ``__core__`` namespace, LangGraph never owns channel merge, and the 

6graph boundary stays flat-dict-in / flat-dict-out. This module supplies the 

7storage, reducer-selection and seeding primitives; runner integration is a 

8separate phase. 

9 

10Stdlib only. Reducer primitives are shared with the aggregate layer through 

11``_reducers``; this module must not import ``langgraph.channels`` or materialize 

12any native LangGraph channel. 

13""" 

14from __future__ import annotations 

15 

16from typing import Any, Callable, Iterable, NamedTuple, Sequence 

17 

18try: 

19 from _channel_reducers import select_channel_reducer 

20 from _reducers import _reduce_add_or_extend, _reduce_last 

21except ImportError: # pragma: no cover — main.py-ajossa (qualified importit) 

22 from langgraph._channel_reducers import select_channel_reducer 

23 from langgraph._reducers import _reduce_add_or_extend, _reduce_last 

24 

25CORE_KEY = "__core__" 

26 

27LAST_VALUE = "LastValue" 

28TOPIC = "Topic" 

29BINARY_OPERATOR_AGGREGATE = "BinaryOperatorAggregate" 

30 

31_KNOWN_TYPES = frozenset({LAST_VALUE, TOPIC, BINARY_OPERATOR_AGGREGATE}) 

32 

33UNSET: Any = object() 

34 

35 

36class ChannelSpec(NamedTuple): 

37 """A declared channel: its flat state key, merge type and declared operator.""" 

38 

39 name: str 

40 type: str 

41 operator: str | None = None 

42 

43 

44def parse_channel_specs(entries: Sequence[dict]) -> list[ChannelSpec]: 

45 """Turn manifest ``spec.channels`` entries into channel specs.""" 

46 return [ 

47 ChannelSpec(entry["name"], entry["type"], entry.get("operator")) 

48 for entry in entries 

49 ] 

50 

51 

52def select_reducer( 

53 channel_type: str, operator: str | None = None 

54) -> Callable[[Any, Any], Any]: 

55 """Select the Core reducer from the channel type and declared operator. 

56 

57 ``LastValue`` and ``Topic`` carry a built-in reducer and ignore 

58 ``operator``; ``BinaryOperatorAggregate`` requires the manifest-declared 

59 operator and resolves it through the name-keyed reducer registry. 

60 """ 

61 _require_known_type(channel_type) 

62 if channel_type == LAST_VALUE: 

63 return _reduce_last 

64 if channel_type == TOPIC: 

65 return _reduce_add_or_extend 

66 if operator is None: 

67 raise ValueError( 

68 f"channel type {BINARY_OPERATOR_AGGREGATE!r} requires an operator" 

69 ) 

70 return select_channel_reducer(operator) 

71 

72 

73def reduce_channel( 

74 channel_type: str, accumulated: Any, write: Any, operator: str | None = None 

75) -> Any: 

76 """Fold one write into a channel's accumulated value (UNSET seeds).""" 

77 if accumulated is UNSET: 

78 return _seed_channel(channel_type, write) 

79 return select_reducer(channel_type, operator)(accumulated, write) 

80 

81 

82def _seed_channel(channel_type: str, write: Any) -> Any: 

83 _require_known_type(channel_type) 

84 if channel_type == TOPIC: 

85 return _reduce_add_or_extend([], write) 

86 return write 

87 

88 

89def _require_known_type(channel_type: str) -> None: 

90 if channel_type not in _KNOWN_TYPES: 

91 raise ValueError(f"unknown channel type: {channel_type!r}") 

92 

93 

94def seed_committed(specs: Iterable[ChannelSpec], initial_state: dict) -> dict: 

95 """Build flat committed state: caller values win, Topic defaults to [].""" 

96 committed = dict(initial_state) 

97 for spec in specs: 

98 if spec.name in committed: 

99 continue 

100 if spec.type == TOPIC: 

101 committed[spec.name] = [] 

102 return committed 

103 

104 

105def _as_channel_spec(name: str, spec: ChannelSpec | str) -> ChannelSpec: 

106 """Normalize a legacy ``name -> type`` entry to a full ``ChannelSpec``.""" 

107 if isinstance(spec, ChannelSpec): 

108 return spec 

109 return ChannelSpec(name, spec) 

110 

111 

112class ChannelWorkingSet: 

113 """Invocation-scoped channel working set with a read-overlay API. 

114 

115 A read observes the committed snapshot overlaid with the invocation's 

116 pending writes, folded in staging order. The committed snapshot itself is 

117 advanced only by ``commit``; the caller's dict is never mutated. 

118 """ 

119 

120 def __init__( 

121 self, committed: dict, channel_specs: dict[str, ChannelSpec | str] 

122 ) -> None: 

123 self._committed = dict(committed) 

124 self._channel_specs = { 

125 name: _as_channel_spec(name, spec) 

126 for name, spec in channel_specs.items() 

127 } 

128 self._pending: dict[str, list[Any]] = {} 

129 

130 def stage(self, name: str, value: Any) -> None: 

131 """Record a pending write for ``name`` in staging order.""" 

132 self._pending.setdefault(name, []).append(value) 

133 

134 def read(self, name: str) -> Any: 

135 """Return the effective value for ``name``: committed plus pending.""" 

136 writes = self._pending.get(name) 

137 if not writes: 

138 return self._committed[name] 

139 return self._fold(name, writes, self._committed.get(name, UNSET)) 

140 

141 def read_view(self) -> dict: 

142 """Return the committed snapshot with pending channels overlaid.""" 

143 view = dict(self._committed) 

144 for name in self._pending: 

145 view[name] = self.read(name) 

146 return view 

147 

148 def fork(self) -> "ChannelWorkingSet": 

149 """Return an isolated working set seeded with this set's read view. 

150 

151 Writes staged on the fork stay invisible to this set until the caller 

152 stages them back, so concurrent workers can accumulate their own 

153 writes without sharing mutable pending state. 

154 """ 

155 return ChannelWorkingSet(self.read_view(), self._channel_specs) 

156 

157 def pending_writes(self) -> list[tuple[str, Any]]: 

158 """Return the staged writes in per-channel staging order.""" 

159 return [ 

160 (name, value) 

161 for name, values in self._pending.items() 

162 for value in values 

163 ] 

164 

165 def commit(self) -> dict: 

166 """Fold pending writes into the committed snapshot and clear them. 

167 

168 Returns the channel-keyed state only, matching the write-commit seam's 

169 shape; non-channel keys survive in the working set's snapshot. 

170 """ 

171 channel_state = { 

172 name: value 

173 for name, value in self._committed.items() 

174 if name in self._channel_specs 

175 } 

176 for name, writes in self._pending.items(): 

177 channel_state[name] = self._fold( 

178 name, writes, channel_state.get(name, UNSET) 

179 ) 

180 self._committed.update(channel_state) 

181 self._pending = {} 

182 return channel_state 

183 

184 def _fold(self, name: str, writes: list[Any], accumulated: Any) -> Any: 

185 spec = self._channel_specs[name] 

186 for write in writes: 

187 accumulated = reduce_channel( 

188 spec.type, accumulated, write, spec.operator 

189 ) 

190 return accumulated 

191 

192 

193def wrap(committed: dict) -> dict: 

194 """Nest flat committed state under the reserved Core namespace.""" 

195 return {CORE_KEY: committed} 

196 

197 

198def unwrap(state: dict) -> dict: 

199 """Return the flat committed state nested under the Core namespace.""" 

200 return state[CORE_KEY]