Coverage for src / lexigram / contracts / ai / runnable.py: 70%

1"""Runnable composition contracts for G-01 parity. 

2 

3Defines the core interfaces for composable runnable components (analogous to 

4LangChain's Runnable interface). 

5""" 

6 

7from __future__ import annotations 

8 

9import asyncio 

10from collections.abc import Callable 

11from dataclasses import dataclass 

12from typing import Any, Protocol, runtime_checkable 

13 

14 

15@dataclass(frozen=True) 

16class RunnableConfigCallbacks: 

17 """Callbacks for runnable lifecycle events. 

18 

19 Attributes: 

20 on_start: Called before runnable execution starts. 

21 on_end: Called after runnable execution completes. 

22 on_error: Called when runnable execution raises an error. 

23 """ 

24 

25 on_start: Callable[[Any], Any] | None = None 

26 on_end: Callable[[Any], Any] | None = None 

27 on_error: Callable[[Exception], Any] | None = None 

28 

29 

30@dataclass(frozen=True) 

31class RunnableConfig: 

32 """Configuration for runnable components. 

33 

34 Attributes: 

35 timeout: Maximum execution time in seconds (default: 30). 

36 max_retries: Maximum retry attempts on failure (default: 3). 

37 callbacks: Lifecycle callbacks for the runnable. 

38 """ 

39 

40 timeout: int = 30 

41 max_retries: int = 3 

42 callbacks: RunnableConfigCallbacks | None = None 

43 

44 

45@runtime_checkable 

46class RunnableProtocol(Protocol): 

47 """Protocol for composable runnable components. 

48 

49 Analogous to LangChain's Runnable interface. Implementations must provide 

50 both sync ``invoke`` and async ``ainvoke`` methods. 

51 """ 

52 

53 def invoke(self, input: Any) -> Any: 

54 """Synchronously process input. 

55 

56 Args: 

57 input: Input to process. 

58 

59 Returns: 

60 Processed output. 

61 """ 

62 ... 

63 

64 async def ainvoke(self, input: Any) -> Any: 

65 """Asynchronously process input. 

66 

67 Args: 

68 input: Input to process. 

69 

70 Returns: 

71 Processed output. 

72 """ 

73 ... 

74 

75 

76class RunnablePipe: 

77 """Compose two runnables in sequence (like LangChain's pipe | operator). 

78 

79 The output of the first runnable becomes the input to the second. 

80 """ 

81 

82 def __init__(self, first: RunnableProtocol, second: RunnableProtocol) -> None: 

83 self.first = first 

84 self.second = second 

85 

86 def invoke(self, input: Any) -> Any: 

87 return self.second.invoke(self.first.invoke(input)) 

88 

89 async def ainvoke(self, input: Any) -> Any: 

90 return await self.second.ainvoke(await self.first.ainvoke(input)) 

91 

92 

93class RunnableParallel: 

94 """Run multiple runnables concurrently (like LangChain's parallel Runnable). 

95 

96 Each runnable receives the same input and results are returned as a dict. 

97 """ 

98 

99 def __init__(self, **runnables: RunnableProtocol) -> None: 

100 self.runnables = runnables 

101 

102 def invoke(self, input: Any) -> dict[str, Any]: 

103 return { 

104 name: runnable.invoke(input) for name, runnable in self.runnables.items() 

105 } 

106 

107 async def ainvoke(self, input: Any) -> dict[str, Any]: 

108 import asyncio 

109 

110 results = await asyncio.gather( 

111 *(runnable.ainvoke(input) for runnable in self.runnables.values()) 

112 ) 

113 return dict(zip(self.runnables.keys(), results, strict=True)) 

114 

115 

116class RunnableLambda: 

117 """Wrap a function as a runnable (like LangChain's RunnableLambda). 

118 

119 Accepts either a sync or async function and wraps it to satisfy 

120 the RunnableProtocol interface. 

121 """ 

122 

123 def __init__(self, func: Callable[[Any], Any]) -> None: 

124 self.func = func 

125 

126 def invoke(self, input: Any) -> Any: 

127 if asyncio.iscoroutinefunction(self.func): 

128 raise TypeError("Use ainvoke for async functions") 

129 return self.func(input) 

130 

131 async def ainvoke(self, input: Any) -> Any: 

132 if asyncio.iscoroutinefunction(self.func): 

133 return await self.func(input) 

134 return self.func(input) 

135 

136 

137class RunnableChain: 

138 """Chain multiple runnables with config support. 

139 

140 Like LangChain's RunnableSequence, this chains multiple runnables 

141 together with optional configuration. 

142 """ 

143 

144 def __init__( 

145 self, 

146 steps: list[RunnableProtocol], 

147 config: RunnableConfig | None = None, 

148 ) -> None: 

149 self.steps = steps 

150 self.config = config or RunnableConfig() 

151 

152 def invoke(self, input: Any) -> Any: 

153 result = input 

154 for step in self.steps: 

155 result = step.invoke(result) 

156 return result 

157 

158 async def ainvoke(self, input: Any) -> Any: 

159 result = input 

160 for step in self.steps: 

161 result = await step.ainvoke(result) 

162 return result 

163 

164 

165class RunnableMap: 

166 """Map input through transforms (like LangChain's RunnableMap). 

167 

168 Accepts either a single transform function or a dict of named transforms. 

169 """ 

170 

171 def __init__( 

172 self, 

173 func_or_funcs: Callable[[Any], Any] | dict[str, Callable[[Any], Any]], 

174 ) -> None: 

175 if isinstance(func_or_funcs, dict): 

176 self.funcs: dict[str, Callable[[Any], Any]] = func_or_funcs 

177 self.single_func: Callable[[Any], Any] | None = None 

178 else: 

179 self.single_func = func_or_funcs 

180 self.funcs = {} 

181 

182 def invoke(self, input: Any) -> Any: 

183 if self.funcs: 

184 return {name: func(input) for name, func in self.funcs.items()} 

185 if self.single_func is not None: 

186 return self.single_func(input) 

187 return input 

188 

189 async def ainvoke(self, input: Any) -> Any: 

190 if self.funcs: 

191 results = await asyncio.gather( 

192 *(func(input) for func in self.funcs.values()) 

193 ) 

194 return dict(zip(self.funcs.keys(), results, strict=True)) 

195 if self.single_func is not None: 

196 if asyncio.iscoroutinefunction(self.single_func): 

197 return await self.single_func(input) 

198 return self.single_func(input) 

199 return input 

200 

201 

202class RunnableGenerator: 

203 """Wrap a generator function for streaming (like LangChain's RunnableGenerator). 

204 

205 Yields chunks for streaming responses. 

206 """ 

207 

208 def __init__(self, generator_func: Callable[[Any], Any]) -> None: 

209 self.generator_func = generator_func 

210 

211 def invoke(self, input: Any) -> list[Any]: 

212 if asyncio.iscoroutinefunction(self.generator_func): 

213 raise TypeError("Use ainvoke for async generators") 

214 return list(self.generator_func(input)) 

215 

216 async def ainvoke(self, input: Any) -> list[Any]: 

217 import inspect 

218 

219 gen = self.generator_func(input) 

220 if inspect.isasyncgenfunction(self.generator_func) or inspect.isasyncgen(gen): 

221 result = [] 

222 async for chunk in gen: 

223 result.append(chunk) 

224 return result 

225 if asyncio.iscoroutinefunction(self.generator_func): 

226 return [await self.generator_func(input)] 

227 return list(gen) 

228 

229 

230class RunnableWithRetry: 

231 """Wrap a runnable with retry logic. 

232 

233 Retries on failure up to max_retries times. 

234 """ 

235 

236 def __init__( 

237 self, 

238 func: Callable[[Any], Any], 

239 max_retries: int = 3, 

240 ) -> None: 

241 self.func = func 

242 self.max_retries = max_retries 

243 

244 def invoke(self, input: Any) -> Any: 

245 last_error: Exception | None = None 

246 for _attempt in range(self.max_retries): 

247 try: 

248 return self.func(input) 

249 except Exception as e: 

250 last_error = e 

251 if last_error is not None: 

252 raise last_error 

253 return input 

254 

255 async def ainvoke(self, input: Any) -> Any: 

256 last_error: Exception | None = None 

257 for _attempt in range(self.max_retries): 

258 try: 

259 if asyncio.iscoroutinefunction(self.func): 

260 return await self.func(input) 

261 return self.func(input) 

262 except Exception as e: 

263 last_error = e 

264 if last_error is not None: 

265 raise last_error 

266 return input 

267 

268 

269__all__ = [ 

270 "RunnableChain", 

271 "RunnableConfig", 

272 "RunnableConfigCallbacks", 

273 "RunnableGenerator", 

274 "RunnableLambda", 

275 "RunnableMap", 

276 "RunnableParallel", 

277 "RunnablePipe", 

278 "RunnableProtocol", 

279 "RunnableWithRetry", 

280]