Coverage for src / lexigram / contracts / ai / callbacks.py: 0%

1"""Callback protocols for AI observability.""" 

2 

3from __future__ import annotations 

4 

5from typing import Any, Protocol, runtime_checkable 

6 

7from lexigram.contracts.ai.llm import ChatMessage, Completion 

8 

9 

10@runtime_checkable 

11class CallbackHandlerProtocol(Protocol): 

12 """Protocol for AI callback handlers. 

13 

14 Defines 12 callback methods for observing LLM, chain, tool, agent, 

15 and retriever operations. Implementations can subscribe to these 

16 events for logging, monitoring, or instrumentation. 

17 """ 

18 

19 async def on_llm_start( 

20 self, 

21 messages: list[ChatMessage], 

22 model: str, 

23 **kwargs: Any, 

24 ) -> None: 

25 """Called when an LLM call starts. 

26 

27 Args: 

28 messages: The messages being sent to the LLM. 

29 model: The model being used. 

30 **kwargs: Additional provider-specific parameters. 

31 """ 

32 ... 

33 

34 async def on_llm_new_token( 

35 self, 

36 token: str, 

37 **kwargs: Any, 

38 ) -> None: 

39 """Called for each new token in a streaming LLM response. 

40 

41 Args: 

42 token: The new token received. 

43 **kwargs: Additional context (e.g., run_id, model). 

44 """ 

45 ... 

46 

47 async def on_llm_end( 

48 self, 

49 response: Completion, 

50 **kwargs: Any, 

51 ) -> None: 

52 """Called when an LLM call completes successfully. 

53 

54 Args: 

55 response: The completion response. 

56 **kwargs: Additional context (e.g., run_id, model). 

57 """ 

58 ... 

59 

60 async def on_llm_error( 

61 self, 

62 error: Exception, 

63 **kwargs: Any, 

64 ) -> None: 

65 """Called when an LLM call fails. 

66 

67 Args: 

68 error: The exception that was raised. 

69 **kwargs: Additional context (e.g., run_id, model, messages). 

70 """ 

71 ... 

72 

73 async def on_chain_start( 

74 self, 

75 name: str, 

76 inputs: dict[str, Any], 

77 **kwargs: Any, 

78 ) -> None: 

79 """Called when a chain/pipeline starts executing. 

80 

81 Args: 

82 name: The name of the chain. 

83 inputs: The input parameters to the chain. 

84 **kwargs: Additional context. 

85 """ 

86 ... 

87 

88 async def on_chain_end( 

89 self, 

90 name: str, 

91 outputs: dict[str, Any], 

92 **kwargs: Any, 

93 ) -> None: 

94 """Called when a chain/pipeline completes. 

95 

96 Args: 

97 name: The name of the chain. 

98 outputs: The output from the chain. 

99 **kwargs: Additional context. 

100 """ 

101 ... 

102 

103 async def on_tool_start( 

104 self, 

105 tool_name: str, 

106 arguments: dict[str, Any], 

107 **kwargs: Any, 

108 ) -> None: 

109 """Called when a tool starts executing. 

110 

111 Args: 

112 tool_name: The name of the tool being invoked. 

113 arguments: The arguments being passed to the tool. 

114 **kwargs: Additional context (e.g., run_id). 

115 """ 

116 ... 

117 

118 async def on_tool_end( 

119 self, 

120 tool_name: str, 

121 result: Any, 

122 **kwargs: Any, 

123 ) -> None: 

124 """Called when a tool finishes executing. 

125 

126 Args: 

127 tool_name: The name of the tool that was invoked. 

128 result: The result from the tool. 

129 **kwargs: Additional context (e.g., run_id). 

130 """ 

131 ... 

132 

133 async def on_agent_action( 

134 self, 

135 action: dict[str, Any], 

136 **kwargs: Any, 

137 ) -> None: 

138 """Called when an agent takes an action. 

139 

140 Args: 

141 action: The action details (e.g., tool call, thought). 

142 **kwargs: Additional context (e.g., run_id). 

143 """ 

144 ... 

145 

146 async def on_agent_finish( 

147 self, 

148 response: dict[str, Any], 

149 **kwargs: Any, 

150 ) -> None: 

151 """Called when an agent finishes executing. 

152 

153 Args: 

154 response: The final response from the agent. 

155 **kwargs: Additional context (e.g., run_id). 

156 """ 

157 ... 

158 

159 async def on_retriever_start( 

160 self, 

161 query: str, 

162 **kwargs: Any, 

163 ) -> None: 

164 """Called when a retriever starts a search. 

165 

166 Args: 

167 query: The search query. 

168 **kwargs: Additional context (e.g., run_id). 

169 """ 

170 ... 

171 

172 async def on_retriever_end( 

173 self, 

174 documents: list[Any], 

175 **kwargs: Any, 

176 ) -> None: 

177 """Called when a retriever completes a search. 

178 

179 Args: 

180 documents: The retrieved documents. 

181 **kwargs: Additional context (e.g., run_id, query). 

182 """ 

183 ... 

184 

185 

186@runtime_checkable 

187class CallbackManagerProtocol(Protocol): 

188 """Protocol for managing callback handlers. 

189 

190 Manages registration and lifecycle of callback handlers, supporting 

191 hierarchical managers (parent-child) for run-scoped event propagation. 

192 """ 

193 

194 def register(self, handler: CallbackHandlerProtocol) -> None: 

195 """Register a callback handler. 

196 

197 Args: 

198 handler: The handler to register. 

199 """ 

200 ... 

201 

202 def unregister(self, handler: CallbackHandlerProtocol) -> None: 

203 """Unregister a callback handler. 

204 

205 Args: 

206 handler: The handler to remove. 

207 """ 

208 ... 

209 

210 def child(self, run_id: str) -> CallbackManagerProtocol: 

211 """Create a child callback manager for a specific run. 

212 

213 The child manager inherits all handlers from the parent and 

214 can be used to scope callbacks to a specific execution run. 

215 

216 Args: 

217 run_id: Unique identifier for this run. 

218 

219 Returns: 

220 A child CallbackManagerProtocol instance. 

221 """ 

222 ... 

223 

224 

225__all__ = [ 

226 "CallbackHandlerProtocol", 

227 "CallbackManagerProtocol", 

228]