Coverage for src / agent_contracts / node.py: 69%

81 statements  

« prev     ^ index     » next       coverage.py v7.13.1, created at 2026-01-09 00:42 +0900

1"""ModularNode and InteractiveNode - Base classes for nodes. 

2 

3All nodes inherit from ModularNode and define a CONTRACT. 

4Contract-based I/O is automatically validated. 

5""" 

6from __future__ import annotations 

7 

8from abc import ABC, abstractmethod 

9from typing import Any, ClassVar, Optional 

10 

11from langchain_core.language_models import BaseChatModel 

12from langchain_core.runnables import RunnableConfig 

13 

14from agent_contracts.contracts import NodeContract, NodeInputs, NodeOutputs 

15from agent_contracts.utils.logging import get_logger 

16 

17 

18class ModularNode(ABC): 

19 """Base class for modular nodes. 

20  

21 All nodes inherit this class and define a CONTRACT class variable. 

22  

23 Example: 

24 class LikeHandlerNode(ModularNode): 

25 CONTRACT = NodeContract( 

26 name="like_handler", 

27 reads=["request", "card", "shopping"], 

28 writes=["card", "shopping", "response"], 

29 ... 

30 ) 

31  

32 async def execute(self, inputs: NodeInputs) -> NodeOutputs: 

33 card = inputs.get_slice("card") 

34 ... 

35 return NodeOutputs(card={...}, response={...}) 

36 """ 

37 

38 # Subclasses must define this 

39 CONTRACT: ClassVar[NodeContract] 

40 

41 def __init__( 

42 self, 

43 llm: BaseChatModel | None = None, 

44 **services: Any, 

45 ): 

46 """Initialize. 

47  

48 Args: 

49 llm: LangChain LLM (required if CONTRACT.requires_llm is True) 

50 **services: Other services (declared in CONTRACT.services) 

51 """ 

52 self.logger = get_logger(self.__class__.__name__) 

53 self.llm = llm 

54 self._validated = False # Deferred validation flag 

55 

56 # Service injection 

57 self._services = services 

58 for service_name in self.CONTRACT.services: 

59 if service_name in services: 

60 setattr(self, service_name, services[service_name]) 

61 

62 def _validate_dependencies(self) -> None: 

63 """Validate declared dependencies from Contract. 

64  

65 Called lazily on first execution to allow subclasses to set 

66 services after calling super().__init__(). 

67 """ 

68 if self._validated: 

69 return 

70 self._validated = True 

71 

72 if self.CONTRACT.requires_llm and self.llm is None: 

73 self.logger.warning( 

74 f"Node {self.CONTRACT.name} requires LLM but none provided" 

75 ) 

76 

77 for service_name in self.CONTRACT.services: 

78 if not hasattr(self, service_name) or getattr(self, service_name) is None: 

79 self.logger.warning( 

80 f"Node {self.CONTRACT.name} requires {service_name} but not provided" 

81 ) 

82 

83 @abstractmethod 

84 async def execute( 

85 self, 

86 inputs: NodeInputs, 

87 config: Optional[RunnableConfig] = None, 

88 ) -> NodeOutputs: 

89 """Execute node's main processing. 

90  

91 Args: 

92 inputs: Input slices per CONTRACT.reads 

93 config: Optional RunnableConfig for LLM tracing. Pass this to 

94 self.llm.ainvoke(..., config=config) for proper tracing. 

95  

96 Returns: 

97 Output slices per CONTRACT.writes 

98 """ 

99 pass 

100 

101 async def __call__( 

102 self, 

103 state: dict, 

104 config: Optional[RunnableConfig] = None, 

105 ) -> dict: 

106 """LangGraph-compatible Callable. 

107  

108 Extracts required slices from State, calls execute, 

109 and converts result to State update format. 

110 """ 

111 # Deferred dependency validation (runs once on first call) 

112 self._validate_dependencies() 

113 

114 # Extract input slices 

115 inputs = self._extract_inputs(state) 

116 

117 # Merge config with node metadata (create new config to avoid mutation) 

118 base_config = config or {} 

119 existing_metadata = base_config.get("metadata", {}) 

120 config = { 

121 **base_config, 

122 "metadata": { 

123 **existing_metadata, 

124 "node_name": self.CONTRACT.name, 

125 "node_supervisor": self.CONTRACT.supervisor, 

126 "node_type": self.__class__.__name__, 

127 }, 

128 } 

129 

130 # Execute 

131 try: 

132 # Execute 

133 outputs = await self.execute(inputs, config=config) 

134 except Exception as e: 

135 self.logger.error(f"Node {self.CONTRACT.name} execution failed: {e}") 

136 raise 

137 

138 # Convert outputs to State update format 

139 return self._convert_outputs(outputs) 

140 

141 def _extract_inputs(self, state: dict) -> NodeInputs: 

142 """Extract required slices from State. 

143  

144 Only extracts slices declared in CONTRACT.reads 

145 and returns them as NodeInputs. 

146 """ 

147 data = {} 

148 for slice_name in self.CONTRACT.reads: 

149 if slice_name == "_internal": 

150 data[slice_name] = state.get("_internal", {}) 

151 else: 

152 data[slice_name] = state.get(slice_name, {}) 

153 

154 return NodeInputs(**data) 

155 

156 def _convert_outputs(self, outputs: NodeOutputs) -> dict: 

157 """Convert NodeOutputs to LangGraph State update format. 

158  

159 Expands from slice format to flat format. 

160 LangGraph expects a flat dict. 

161 """ 

162 result = {} 

163 for slice_name, slice_data in outputs.to_state_updates().items(): 

164 if isinstance(slice_data, dict): 

165 result[slice_name] = slice_data 

166 return result 

167 

168 # ========================================================================= 

169 # Helper Methods (for subclasses) 

170 # ========================================================================= 

171 

172 def get_request_param(self, inputs: NodeInputs, key: str, default: Any = None) -> Any: 

173 """Get request parameter.""" 

174 request = inputs.get_slice("request") 

175 params = request.get("params") or {} 

176 return params.get(key, default) 

177 

178 def build_error_response(self, message: str, code: str) -> NodeOutputs: 

179 """Build error response.""" 

180 return NodeOutputs( 

181 response={ 

182 "response_type": "error", 

183 "response_data": {"message": message, "code": code}, 

184 } 

185 ) 

186 

187 

188class InteractiveNode(ModularNode): 

189 """Base class for conversational nodes. 

190  

191 Provides standard flow for: 

192 1. prepare_context: Prepare execution context 

193 2. process_answer: Answer processing (if previous question exists) 

194 3. check_completion: Completion check 

195 4. generate_question: Next question generation (if not complete) 

196  

197 Subclasses should implement: 

198 - prepare_context(inputs) -> Any 

199 - check_completion(context, inputs) -> bool 

200 - process_answer(context, inputs) -> bool 

201 - generate_question(context, inputs) -> NodeOutputs 

202 """ 

203 

204 @abstractmethod 

205 def prepare_context(self, inputs: NodeInputs) -> Any: 

206 """Prepare execution context. 

207  

208 Extract needed data from NodeInputs and convert to 

209 easy-to-use object (Pydantic model, etc.). 

210 """ 

211 pass 

212 

213 @abstractmethod 

214 def check_completion(self, context: Any, inputs: NodeInputs) -> bool: 

215 """Check task completion.""" 

216 pass 

217 

218 @abstractmethod 

219 async def process_answer( 

220 self, 

221 context: Any, 

222 inputs: NodeInputs, 

223 config: RunnableConfig | None = None 

224 ) -> bool: 

225 """Process user answer. 

226  

227 Returns: 

228 bool: True if answer was processed and state updated 

229 """ 

230 pass 

231 

232 @abstractmethod 

233 async def generate_question( 

234 self, 

235 context: Any, 

236 inputs: NodeInputs, 

237 config: RunnableConfig | None = None 

238 ) -> NodeOutputs: 

239 """Generate and return next question.""" 

240 pass 

241 

242 async def create_completion_output( 

243 self, 

244 context: Any, 

245 inputs: NodeInputs, 

246 config: RunnableConfig | None = None 

247 ) -> NodeOutputs: 

248 """Create output for completion (default: done flag).""" 

249 return NodeOutputs(_internal={"decision": "done"}) 

250 

251 async def execute( 

252 self, 

253 inputs: NodeInputs, 

254 config: Optional[RunnableConfig] = None, 

255 ) -> NodeOutputs: 

256 """Standard execution flow.""" 

257 

258 # 0. Prepare context 

259 context = self.prepare_context(inputs) 

260 

261 # 1. Process answer 

262 await self.process_answer(context, inputs, config=config) 

263 

264 # 2. Check completion 

265 if self.check_completion(context, inputs): 

266 return await self.create_completion_output(context, inputs, config=config) 

267 

268 # 3. Generate question 

269 return await self.generate_question(context, inputs, config=config)