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 _reducers import ( 

20 _is_collection_map, 

21 _reduce_add_or_extend, 

22 _reduce_dict_merge, 

23 _reduce_last, 

24 _reduce_union_by_key, 

25 _require_mapping, 

26 ) 

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

28 from langgraph._reducers import ( 

29 _is_collection_map, 

30 _reduce_add_or_extend, 

31 _reduce_dict_merge, 

32 _reduce_last, 

33 _reduce_union_by_key, 

34 _require_mapping, 

35 ) 

36 

37CORE_KEY = "__core__" 

38 

39LAST_VALUE = "LastValue" 

40TOPIC = "Topic" 

41BINARY_OPERATOR_AGGREGATE = "BinaryOperatorAggregate" 

42 

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

44 

45UNSET: Any = object() 

46 

47 

48class ChannelSpec(NamedTuple): 

49 """A declared channel: its flat state key and its merge type.""" 

50 

51 name: str 

52 type: str 

53 

54 

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

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

57 return [ChannelSpec(entry["name"], entry["type"]) for entry in entries] 

58 

59 

60def select_reducer( 

61 channel_type: str, accumulated: Any 

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

63 """Select the Core reducer from the channel type and accumulated shape.""" 

64 _require_known_type(channel_type) 

65 if channel_type == LAST_VALUE: 

66 return _reduce_last 

67 if channel_type == TOPIC: 

68 return _reduce_add_or_extend 

69 return _select_binop_reducer(accumulated) 

70 

71 

72def _select_binop_reducer(accumulated: Any) -> Callable[[Any, Any], Any]: 

73 if isinstance(accumulated, list): 

74 return _reduce_add_or_extend 

75 if _is_collection_map(accumulated): 

76 return _reduce_union_by_key 

77 if isinstance(accumulated, dict): 

78 return _reduce_dict_merge 

79 return _reduce_last 

80 

81 

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

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

84 if accumulated is UNSET: 

85 return _seed_channel(channel_type, write) 

86 if channel_type == BINARY_OPERATOR_AGGREGATE and isinstance(accumulated, dict): 

87 _require_mapping(write) 

88 return select_reducer(channel_type, accumulated)(accumulated, write) 

89 

90 

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

92 _require_known_type(channel_type) 

93 if channel_type == TOPIC: 

94 return _reduce_add_or_extend([], write) 

95 return write 

96 

97 

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

99 if channel_type not in _KNOWN_TYPES: 

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

101 

102 

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

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

105 committed = dict(initial_state) 

106 for spec in specs: 

107 if spec.name in committed: 

108 continue 

109 if spec.type == TOPIC: 

110 committed[spec.name] = [] 

111 return committed 

112 

113 

114class ChannelWorkingSet: 

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

116 

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

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

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

120 """ 

121 

122 def __init__(self, committed: dict, channel_types: dict[str, str]) -> None: 

123 self._committed = dict(committed) 

124 self._channel_types = dict(channel_types) 

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

126 

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

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

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

130 

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

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

133 writes = self._pending.get(name) 

134 if not writes: 

135 return self._committed[name] 

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

137 

138 def read_view(self) -> dict: 

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

140 view = dict(self._committed) 

141 for name in self._pending: 

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

143 return view 

144 

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

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

147 

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

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

150 writes without sharing mutable pending state. 

151 """ 

152 return ChannelWorkingSet(self.read_view(), self._channel_types) 

153 

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

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

156 return [ 

157 (name, value) 

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

159 for value in values 

160 ] 

161 

162 def commit(self) -> dict: 

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

164 

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

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

167 """ 

168 channel_state = { 

169 name: value 

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

171 if name in self._channel_types 

172 } 

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

174 channel_state[name] = self._fold( 

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

176 ) 

177 self._committed.update(channel_state) 

178 self._pending = {} 

179 return channel_state 

180 

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

182 for write in writes: 

183 accumulated = reduce_channel( 

184 self._channel_types[name], accumulated, write 

185 ) 

186 return accumulated 

187 

188 

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

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

191 return {CORE_KEY: committed} 

192 

193 

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

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

196 return state[CORE_KEY]