Coverage for haystack/dataclasses/streaming_chunk.py: 99%

1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai> 

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5import inspect 

6from collections.abc import Awaitable, Callable 

7from dataclasses import asdict, dataclass, field 

8from typing import Any, Literal, overload 

9 

10from haystack import logging 

11from haystack.core.component import Component 

12from haystack.dataclasses.chat_message import ReasoningContent, ToolCallResult 

13from haystack.utils.dataclasses import _warn_on_inplace_mutation 

14 

15logger = logging.getLogger(__name__) 

16 

17# Type alias for standard finish_reason values following OpenAI's convention 

18# plus Haystack-specific value ("tool_call_results") 

19FinishReason = Literal["stop", "length", "tool_calls", "content_filter", "tool_call_results"] 

20 

21 

22@_warn_on_inplace_mutation 

23@dataclass 

24class ToolCallDelta: 

25 """ 

26 Represents a Tool call prepared by the model, usually contained in an assistant message. 

27 

28 :param index: The index of the Tool call in the list of Tool calls. 

29 :param tool_name: The name of the Tool to call. 

30 :param arguments: Either the full arguments in JSON format or a delta of the arguments. 

31 :param id: The ID of the Tool call. 

32 :param extra: Dictionary of extra information about the Tool call. Use to store provider-specific 

33 information. To avoid serialization issues, values should be JSON serializable. 

34 """ 

35 

36 index: int 

37 tool_name: str | None = field(default=None) 

38 arguments: str | None = field(default=None) 

39 id: str | None = field(default=None) 

40 extra: dict[str, Any] | None = field(default=None) 

41 

42 def to_dict(self) -> dict[str, Any]: 

43 """ 

44 Returns a dictionary representation of the ToolCallDelta. 

45 

46 :returns: A dictionary with keys 'index', 'tool_name', 'arguments', 'id', and 'extra'. 

47 """ 

48 return asdict(self) 

49 

50 @classmethod 

51 def from_dict(cls, data: dict[str, Any]) -> "ToolCallDelta": 

52 """ 

53 Creates a ToolCallDelta from a serialized representation. 

54 

55 :param data: Dictionary containing ToolCallDelta's attributes. 

56 :returns: A ToolCallDelta instance. 

57 """ 

58 return ToolCallDelta(**data) 

59 

60 

61@_warn_on_inplace_mutation 

62@dataclass 

63class ComponentInfo: 

64 """ 

65 The `ComponentInfo` class encapsulates information about a component. 

66 

67 :param type: The type of the component. 

68 :param name: The name of the component assigned when adding it to a pipeline. 

69 

70 """ 

71 

72 type: str 

73 name: str | None = field(default=None) 

74 

75 @classmethod 

76 def from_component(cls, component: Component) -> "ComponentInfo": 

77 """ 

78 Create a `ComponentInfo` object from a `Component` instance. 

79 

80 :param component: 

81 The `Component` instance. 

82 :returns: 

83 The `ComponentInfo` object with the type and name of the given component. 

84 """ 

85 component_type = f"{component.__class__.__module__}.{component.__class__.__name__}" 

86 component_name = getattr(component, "__component_name__", None) 

87 return cls(type=component_type, name=component_name) 

88 

89 def to_dict(self) -> dict[str, Any]: 

90 """ 

91 Returns a dictionary representation of ComponentInfo. 

92 

93 :returns: A dictionary with keys 'type' and 'name'. 

94 """ 

95 return asdict(self) 

96 

97 @classmethod 

98 def from_dict(cls, data: dict[str, Any]) -> "ComponentInfo": 

99 """ 

100 Creates a ComponentInfo from a serialized representation. 

101 

102 :param data: Dictionary containing ComponentInfo's attributes. 

103 :returns: A ComponentInfo instance. 

104 """ 

105 return ComponentInfo(**data) 

106 

107 

108@_warn_on_inplace_mutation 

109@dataclass 

110class StreamingChunk: 

111 """ 

112 The `StreamingChunk` class encapsulates a segment of streamed content along with associated metadata. 

113 

114 This structure facilitates the handling and processing of streamed data in a systematic manner. 

115 

116 :param content: The content of the message chunk as a string. 

117 :param meta: A dictionary containing metadata related to the message chunk. 

118 :param component_info: A `ComponentInfo` object containing information about the component that generated the chunk, 

119 such as the component name and type. 

120 :param index: An optional integer index representing which content block this chunk belongs to. 

121 :param tool_calls: An optional list of ToolCallDelta object representing a tool call associated with the message 

122 chunk. 

123 :param tool_call_result: An optional ToolCallResult object representing the result of a tool call. 

124 :param start: A boolean indicating whether this chunk marks the start of a content block. 

125 :param finish_reason: An optional value indicating the reason the generation finished. 

126 Standard values follow OpenAI's convention: "stop", "length", "tool_calls", "content_filter", 

127 plus Haystack-specific value "tool_call_results". 

128 :param reasoning: An optional ReasoningContent object representing the reasoning content associated 

129 with the message chunk. 

130 """ 

131 

132 content: str 

133 meta: dict[str, Any] = field(default_factory=dict, hash=False) 

134 component_info: ComponentInfo | None = field(default=None) 

135 index: int | None = field(default=None) 

136 tool_calls: list[ToolCallDelta] | None = field(default=None) 

137 tool_call_result: ToolCallResult | None = field(default=None) 

138 start: bool = field(default=False) 

139 finish_reason: FinishReason | None = field(default=None) 

140 reasoning: ReasoningContent | None = field(default=None) 

141 

142 def __post_init__(self) -> None: 

143 fields_set = sum(bool(x) for x in (self.content, self.tool_calls, self.tool_call_result, self.reasoning)) 

144 if fields_set > 1: 

145 raise ValueError( 

146 "Only one of `content`, `tool_call`, `tool_call_result` or `reasoning` may be set in a StreamingChunk. " 

147 f"Got content: '{self.content}', tool_call: '{self.tool_calls}', " 

148 f"tool_call_result: '{self.tool_call_result}', reasoning: '{self.reasoning}'." 

149 ) 

150 

151 # NOTE: We don't enforce this for self.content otherwise it would be a breaking change 

152 if (self.tool_calls or self.tool_call_result or self.reasoning) and self.index is None: 

153 raise ValueError("If `tool_call`, `tool_call_result` or `reasoning` is set, `index` must also be set.") 

154 

155 def to_dict(self) -> dict[str, Any]: 

156 """ 

157 Returns a dictionary representation of the StreamingChunk. 

158 

159 :returns: Serialized dictionary representation of the calling object. 

160 """ 

161 return { 

162 "content": self.content, 

163 "meta": self.meta, 

164 "component_info": self.component_info.to_dict() if self.component_info else None, 

165 "index": self.index, 

166 "tool_calls": [tc.to_dict() for tc in self.tool_calls] if self.tool_calls else None, 

167 "tool_call_result": self.tool_call_result.to_dict() if self.tool_call_result else None, 

168 "start": self.start, 

169 "finish_reason": self.finish_reason, 

170 "reasoning": self.reasoning.to_dict() if self.reasoning else None, 

171 } 

172 

173 @classmethod 

174 def from_dict(cls, data: dict[str, Any]) -> "StreamingChunk": 

175 """ 

176 Creates a deserialized StreamingChunk instance from a serialized representation. 

177 

178 :param data: Dictionary containing the StreamingChunk's attributes. 

179 :returns: A StreamingChunk instance. 

180 """ 

181 if "content" not in data: 

182 raise ValueError("Missing required field `content` in StreamingChunk deserialization.") 

183 

184 return StreamingChunk( 

185 content=data["content"], 

186 meta=data.get("meta", {}), 

187 component_info=ComponentInfo.from_dict(data["component_info"]) if data.get("component_info") else None, 

188 index=data.get("index"), 

189 tool_calls=[ToolCallDelta.from_dict(tc) for tc in data["tool_calls"]] if data.get("tool_calls") else None, 

190 tool_call_result=ToolCallResult.from_dict(data["tool_call_result"]) 

191 if data.get("tool_call_result") 

192 else None, 

193 start=data.get("start", False), 

194 finish_reason=data.get("finish_reason"), 

195 reasoning=ReasoningContent.from_dict(data["reasoning"]) if data.get("reasoning") else None, 

196 ) 

197 

198 

199SyncStreamingCallbackT = Callable[[StreamingChunk], None] 

200AsyncStreamingCallbackT = Callable[[StreamingChunk], Awaitable[None]] 

201 

202StreamingCallbackT = SyncStreamingCallbackT | AsyncStreamingCallbackT 

203 

204 

205def _is_callable_async_compatible(func: Callable) -> bool: 

206 """ 

207 Returns if the given callable is usable inside a component's `run_async` method. 

208 

209 :param func: 

210 The callable to check. 

211 :returns: 

212 True if the callable is compatible, False otherwise. 

213 """ 

214 return inspect.iscoroutinefunction(func) 

215 

216 

217@overload 

218def select_streaming_callback( 

219 init_callback: StreamingCallbackT | None, 

220 runtime_callback: StreamingCallbackT | None, 

221 requires_async: Literal[False], 

222) -> SyncStreamingCallbackT | None: ... 

223@overload 

224def select_streaming_callback( 

225 init_callback: StreamingCallbackT | None, runtime_callback: StreamingCallbackT | None, requires_async: Literal[True] 

226) -> StreamingCallbackT | None: ... 

227 

228 

229def select_streaming_callback( 

230 init_callback: StreamingCallbackT | None, runtime_callback: StreamingCallbackT | None, requires_async: bool 

231) -> StreamingCallbackT | None: 

232 """ 

233 Picks the correct streaming callback given an optional initial and runtime callback. 

234 

235 The runtime callback takes precedence over the initial callback. 

236 

237 In an async context (`requires_async=True`), a sync callback is accepted but emits a warning: it will run inline on 

238 the event loop and may block it. In a sync context (`requires_async=False`), an async callback is rejected because 

239 there is no way to await it. 

240 

241 :param init_callback: 

242 The initial callback. 

243 :param runtime_callback: 

244 The runtime callback. 

245 :param requires_async: 

246 Whether the selected callback will be invoked from an async context. 

247 :returns: 

248 The selected callback. 

249 """ 

250 if init_callback is not None: 

251 if requires_async and not _is_callable_async_compatible(init_callback): 

252 logger.warning( 

253 "A sync streaming callback was provided at initialization for use in an async context. " 

254 "It will run synchronously on the event loop and may block it." 

255 ) 

256 if not requires_async and _is_callable_async_compatible(init_callback): 

257 raise ValueError("The init callback cannot be a coroutine.") 

258 

259 if runtime_callback is not None: 

260 if requires_async and not _is_callable_async_compatible(runtime_callback): 

261 logger.warning( 

262 "A sync streaming callback was provided at runtime for use in an async context. " 

263 "It will run synchronously on the event loop and may block it." 

264 ) 

265 if not requires_async and _is_callable_async_compatible(runtime_callback): 

266 raise ValueError("The runtime callback cannot be a coroutine.") 

267 

268 return runtime_callback or init_callback 

269 

270 

271async def _invoke_streaming_callback(callback: StreamingCallbackT, chunk: StreamingChunk) -> None: 

272 """ 

273 Invokes a streaming callback in an async context, handling both sync and async callbacks. 

274 

275 :param callback: The streaming callback to invoke. 

276 :param chunk: The streaming chunk to pass to the callback. 

277 """ 

278 result = callback(chunk) 

279 if inspect.isawaitable(result): 

280 await result