haystack-experimental 0.14.2__tar.gz → 0.14.3__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (56) hide show
  1. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/PKG-INFO +1 -1
  2. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/components/agents/agent.py +61 -15
  3. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/components/preprocessors/__init__.py +2 -0
  4. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/components/preprocessors/md_header_level_inferrer.py +2 -2
  5. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/components/query/query_expander.py +13 -18
  6. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/components/retrievers/multi_query_embedding_retriever.py +8 -15
  7. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/components/retrievers/multi_query_text_retriever.py +3 -11
  8. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/core/pipeline/breakpoint.py +5 -3
  9. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/.gitignore +0 -0
  10. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/LICENSE +0 -0
  11. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/LICENSE-MIT.txt +0 -0
  12. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/README.md +0 -0
  13. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/__init__.py +0 -0
  14. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/chat_message_stores/__init__.py +0 -0
  15. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/chat_message_stores/in_memory.py +0 -0
  16. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/chat_message_stores/types.py +0 -0
  17. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/components/__init__.py +0 -0
  18. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/components/agents/__init__.py +0 -0
  19. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/components/agents/human_in_the_loop/__init__.py +0 -0
  20. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/components/agents/human_in_the_loop/breakpoint.py +0 -0
  21. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/components/agents/human_in_the_loop/dataclasses.py +0 -0
  22. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/components/agents/human_in_the_loop/errors.py +0 -0
  23. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/components/agents/human_in_the_loop/policies.py +0 -0
  24. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/components/agents/human_in_the_loop/strategies.py +0 -0
  25. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/components/agents/human_in_the_loop/types.py +0 -0
  26. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/components/agents/human_in_the_loop/user_interfaces.py +0 -0
  27. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/components/embedders/__init__.py +0 -0
  28. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/components/embedders/types/__init__.py +0 -0
  29. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/components/embedders/types/protocol.py +0 -0
  30. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/components/generators/__init__.py +0 -0
  31. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/components/generators/chat/__init__.py +0 -0
  32. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/components/generators/chat/openai.py +0 -0
  33. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/components/preprocessors/embedding_based_document_splitter.py +0 -0
  34. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/components/query/__init__.py +0 -0
  35. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/components/retrievers/__init__.py +0 -0
  36. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/components/retrievers/chat_message_retriever.py +0 -0
  37. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/components/retrievers/types/__init__.py +0 -0
  38. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/components/retrievers/types/protocol.py +0 -0
  39. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/components/summarizers/__init__.py +0 -0
  40. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/components/summarizers/llm_summarizer.py +0 -0
  41. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/components/writers/__init__.py +0 -0
  42. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/components/writers/chat_message_writer.py +0 -0
  43. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/core/__init__.py +0 -0
  44. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/core/pipeline/__init__.py +0 -0
  45. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/dataclasses/__init__.py +0 -0
  46. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/dataclasses/breakpoints.py +0 -0
  47. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/super_components/__init__.py +0 -0
  48. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/super_components/indexers/__init__.py +0 -0
  49. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/super_components/indexers/sentence_transformers_document_indexer.py +0 -0
  50. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/utils/__init__.py +0 -0
  51. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/utils/hallucination_risk_calculator/__init__.py +0 -0
  52. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/utils/hallucination_risk_calculator/core_math.py +0 -0
  53. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/utils/hallucination_risk_calculator/dataclasses.py +0 -0
  54. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/utils/hallucination_risk_calculator/openai_planner.py +0 -0
  55. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/haystack_experimental/utils/hallucination_risk_calculator/skeletonization.py +0 -0
  56. {haystack_experimental-0.14.2 → haystack_experimental-0.14.3}/pyproject.toml +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: haystack-experimental
3
- Version: 0.14.2
3
+ Version: 0.14.3
4
4
  Summary: Experimental components and features for the Haystack LLM framework.
5
5
  Project-URL: CI: GitHub, https://github.com/deepset-ai/haystack-experimental/actions
6
6
  Project-URL: GitHub: issues, https://github.com/deepset-ai/haystack-experimental/issues
@@ -172,6 +172,7 @@ class Agent(HaystackAgent):
172
172
  requires_async: bool,
173
173
  *,
174
174
  system_prompt: Optional[str] = None,
175
+ generation_kwargs: Optional[dict[str, Any]] = None,
175
176
  tools: Optional[Union[ToolsType, list[str]]] = None,
176
177
  **kwargs: dict[str, Any],
177
178
  ) -> _ExecutionContext:
@@ -186,14 +187,28 @@ class Agent(HaystackAgent):
186
187
  When passing tool names, tools are selected from the Agent's originally configured tools.
187
188
  :param kwargs: Additional data to pass to the State used by the Agent.
188
189
  """
189
- exe_context = super(Agent, self)._initialize_fresh_execution(
190
- messages=messages,
191
- streaming_callback=streaming_callback,
192
- requires_async=requires_async,
193
- system_prompt=system_prompt,
194
- tools=tools,
195
- **kwargs,
196
- )
190
+ # The PR https://github.com/deepset-ai/haystack/pull/9616 added the generation_kwargs parameter to
191
+ # _initialize_fresh_execution. This change has been released in Haystack 2.20.0.
192
+ # To maintain compatibility with Haystack 2.19 we check the number of parameters and call accordingly.
193
+ if inspect.signature(super(Agent, self)._initialize_fresh_execution).parameters.get("generation_kwargs"):
194
+ exe_context = super(Agent, self)._initialize_fresh_execution(
195
+ messages=messages,
196
+ streaming_callback=streaming_callback,
197
+ requires_async=requires_async,
198
+ system_prompt=system_prompt,
199
+ generation_kwargs=generation_kwargs,
200
+ tools=tools,
201
+ **kwargs,
202
+ )
203
+ else:
204
+ exe_context = super(Agent, self)._initialize_fresh_execution(
205
+ messages=messages,
206
+ streaming_callback=streaming_callback,
207
+ requires_async=requires_async,
208
+ system_prompt=system_prompt,
209
+ tools=tools,
210
+ **kwargs,
211
+ )
197
212
  # NOTE: 1st difference with parent method to add this to tool_invoker_inputs
198
213
  if self._tool_invoker:
199
214
  exe_context.tool_invoker_inputs["enable_streaming_callback_passthrough"] = (
@@ -213,6 +228,7 @@ class Agent(HaystackAgent):
213
228
  streaming_callback: Optional[StreamingCallbackT],
214
229
  requires_async: bool,
215
230
  *,
231
+ generation_kwargs: Optional[dict[str, Any]] = None,
216
232
  tools: Optional[Union[ToolsType, list[str]]] = None,
217
233
  ) -> _ExecutionContext:
218
234
  """
@@ -221,12 +237,26 @@ class Agent(HaystackAgent):
221
237
  :param snapshot: An AgentSnapshot containing the state of a previously saved agent execution.
222
238
  :param streaming_callback: Optional callback for streaming responses.
223
239
  :param requires_async: Whether the agent run requires asynchronous execution.
240
+ :param generation_kwargs: Additional keyword arguments for chat generator. These parameters will
241
+ override the parameters passed during component initialization.
224
242
  :param tools: Optional list of Tool objects, a Toolset, or list of tool names to use for this run.
225
243
  When passing tool names, tools are selected from the Agent's originally configured tools.
226
244
  """
227
- exe_context = super(Agent, self)._initialize_from_snapshot(
228
- snapshot=snapshot, streaming_callback=streaming_callback, requires_async=requires_async, tools=tools
229
- )
245
+ # The PR https://github.com/deepset-ai/haystack/pull/9616 added the generation_kwargs parameter to
246
+ # _initialize_from_snapshot. This change has been released in Haystack 2.20.0.
247
+ # To maintain compatibility with Haystack 2.19 we check the number of parameters and call accordingly.
248
+ if inspect.signature(super(Agent, self)._initialize_from_snapshot).parameters.get("generation_kwargs"):
249
+ exe_context = super(Agent, self)._initialize_from_snapshot(
250
+ snapshot=snapshot,
251
+ streaming_callback=streaming_callback,
252
+ requires_async=requires_async,
253
+ generation_kwargs=generation_kwargs,
254
+ tools=tools,
255
+ )
256
+ else:
257
+ exe_context = super(Agent, self)._initialize_from_snapshot(
258
+ snapshot=snapshot, streaming_callback=streaming_callback, requires_async=requires_async, tools=tools
259
+ )
230
260
  # NOTE: 1st difference with parent method to add this to tool_invoker_inputs
231
261
  if self._tool_invoker:
232
262
  exe_context.tool_invoker_inputs["enable_streaming_callback_passthrough"] = (
@@ -248,6 +278,7 @@ class Agent(HaystackAgent):
248
278
  messages: list[ChatMessage],
249
279
  streaming_callback: Optional[StreamingCallbackT] = None,
250
280
  *,
281
+ generation_kwargs: Optional[dict[str, Any]] = None,
251
282
  break_point: Optional[AgentBreakpoint] = None,
252
283
  snapshot: Optional[AgentSnapshot] = None, # type: ignore[override]
253
284
  system_prompt: Optional[str] = None,
@@ -260,6 +291,8 @@ class Agent(HaystackAgent):
260
291
  :param messages: List of Haystack ChatMessage objects to process.
261
292
  :param streaming_callback: A callback that will be invoked when a response is streamed from the LLM.
262
293
  The same callback can be configured to emit tool results when a tool is called.
294
+ :param generation_kwargs: Additional keyword arguments for LLM. These parameters will
295
+ override the parameters passed during component initialization.
263
296
  :param break_point: An AgentBreakpoint, can be a Breakpoint for the "chat_generator" or a ToolBreakpoint
264
297
  for "tool_invoker".
265
298
  :param snapshot: A dictionary containing a snapshot of a previously saved agent execution. The snapshot contains
@@ -290,13 +323,17 @@ class Agent(HaystackAgent):
290
323
  # _runtime_checks. This change will be released in Haystack 2.20.0.
291
324
  # To maintain compatibility with Haystack 2.19 we check the number of parameters and call accordingly.
292
325
  if len(inspect.signature(self._runtime_checks).parameters) == 2:
293
- self._runtime_checks(break_point, snapshot)
326
+ self._runtime_checks(break_point, snapshot) # type: ignore[call-arg] # pylint: disable=too-many-function-args
294
327
  else:
295
328
  self._runtime_checks(break_point) # type: ignore[call-arg] # pylint: disable=no-value-for-parameter
296
329
 
297
330
  if snapshot:
298
331
  exe_context = self._initialize_from_snapshot(
299
- snapshot=snapshot, streaming_callback=streaming_callback, requires_async=False, tools=tools
332
+ snapshot=snapshot,
333
+ streaming_callback=streaming_callback,
334
+ requires_async=False,
335
+ generation_kwargs=generation_kwargs,
336
+ tools=tools,
300
337
  )
301
338
  else:
302
339
  exe_context = self._initialize_fresh_execution(
@@ -304,6 +341,7 @@ class Agent(HaystackAgent):
304
341
  streaming_callback=streaming_callback,
305
342
  requires_async=False,
306
343
  system_prompt=system_prompt,
344
+ generation_kwargs=generation_kwargs,
307
345
  tools=tools,
308
346
  **kwargs,
309
347
  )
@@ -438,6 +476,7 @@ class Agent(HaystackAgent):
438
476
  messages: list[ChatMessage],
439
477
  streaming_callback: Optional[StreamingCallbackT] = None,
440
478
  *,
479
+ generation_kwargs: Optional[dict[str, Any]] = None,
441
480
  break_point: Optional[AgentBreakpoint] = None,
442
481
  snapshot: Optional[AgentSnapshot] = None, # type: ignore[override]
443
482
  system_prompt: Optional[str] = None,
@@ -454,6 +493,8 @@ class Agent(HaystackAgent):
454
493
  :param messages: List of Haystack ChatMessage objects to process.
455
494
  :param streaming_callback: An asynchronous callback that will be invoked when a response is streamed from the
456
495
  LLM. The same callback can be configured to emit tool results when a tool is called.
496
+ :param generation_kwargs: Additional keyword arguments for LLM. These parameters will
497
+ override the parameters passed during component initialization.
457
498
  :param break_point: An AgentBreakpoint, can be a Breakpoint for the "chat_generator" or a ToolBreakpoint
458
499
  for "tool_invoker".
459
500
  :param snapshot: A dictionary containing a snapshot of a previously saved agent execution. The snapshot contains
@@ -483,13 +524,17 @@ class Agent(HaystackAgent):
483
524
  # _runtime_checks. This change will be released in Haystack 2.20.0.
484
525
  # To maintain compatibility with Haystack 2.19 we check the number of parameters and call accordingly.
485
526
  if len(inspect.signature(self._runtime_checks).parameters) == 2:
486
- self._runtime_checks(break_point, snapshot)
527
+ self._runtime_checks(break_point, snapshot) # type: ignore[call-arg] # pylint: disable=too-many-function-args
487
528
  else:
488
529
  self._runtime_checks(break_point) # type: ignore[call-arg] # pylint: disable=no-value-for-parameter
489
530
 
490
531
  if snapshot:
491
532
  exe_context = self._initialize_from_snapshot(
492
- snapshot=snapshot, streaming_callback=streaming_callback, requires_async=True, tools=tools
533
+ snapshot=snapshot,
534
+ streaming_callback=streaming_callback,
535
+ requires_async=True,
536
+ generation_kwargs=generation_kwargs,
537
+ tools=tools,
493
538
  )
494
539
  else:
495
540
  exe_context = self._initialize_fresh_execution(
@@ -497,6 +542,7 @@ class Agent(HaystackAgent):
497
542
  streaming_callback=streaming_callback,
498
543
  requires_async=True,
499
544
  system_prompt=system_prompt,
545
+ generation_kwargs=generation_kwargs,
500
546
  tools=tools,
501
547
  **kwargs,
502
548
  )
@@ -9,10 +9,12 @@ from lazy_imports import LazyImporter
9
9
 
10
10
  _import_structure = {
11
11
  "embedding_based_document_splitter": ["EmbeddingBasedDocumentSplitter"],
12
+ "md_header_level_inferrer": ["MarkdownHeaderLevelInferrer"],
12
13
  }
13
14
 
14
15
  if TYPE_CHECKING:
15
16
  from .embedding_based_document_splitter import EmbeddingBasedDocumentSplitter
17
+ from .md_header_level_inferrer import MarkdownHeaderLevelInferrer
16
18
 
17
19
  else:
18
20
  sys.modules[__name__] = LazyImporter(name=__name__, module_file=__file__, import_structure=_import_structure)
@@ -24,7 +24,7 @@ class MarkdownHeaderLevelInferrer:
24
24
  from haystack_experimental.components.preprocessors import MarkdownHeaderLevelInferrer
25
25
 
26
26
  # Create a document with uniform header levels
27
- text = "## Title\nSome content\n## Section\nMore content\n## Subsection\nFinal content"
27
+ text = "## Title\n## Subheader\nSection\n## Subheader\nMore Content"
28
28
  doc = Document(content=text)
29
29
 
30
30
  # Initialize the inferrer and process the document
@@ -33,7 +33,7 @@ class MarkdownHeaderLevelInferrer:
33
33
 
34
34
  # The headers are now normalized with proper hierarchy
35
35
  print(result["documents"][0].content)
36
- > # Title\nSome content\n## Section\nMore content\n### Subsection\nFinal content
36
+ > # Title\n## Subheader\nSection\n## Subheader\nMore Content
37
37
  ```
38
38
  """
39
39
 
@@ -3,7 +3,7 @@
3
3
  # SPDX-License-Identifier: Apache-2.0
4
4
 
5
5
  import json
6
- from typing import Any, Dict, List, Optional
6
+ from typing import Any, Optional
7
7
 
8
8
  from haystack import default_from_dict, default_to_dict, logging
9
9
  from haystack.components.builders.prompt_builder import PromptBuilder
@@ -91,7 +91,7 @@ class QueryExpander:
91
91
  prompt_template: Optional[str] = None,
92
92
  n_expansions: int = 4,
93
93
  include_original_query: bool = True,
94
- ):
94
+ ) -> None:
95
95
  """
96
96
  Initialize the QueryExpander component.
97
97
 
@@ -134,7 +134,6 @@ class QueryExpander:
134
134
  self.chat_generator = chat_generator
135
135
 
136
136
  self._is_warmed_up = False
137
- self._supports_warm_up = hasattr(self.chat_generator, "warm_up")
138
137
  self.prompt_template = prompt_template or DEFAULT_PROMPT_TEMPLATE
139
138
 
140
139
  # Check if required variables are present in the template
@@ -153,7 +152,7 @@ class QueryExpander:
153
152
  required_variables=["n_expansions", "query"],
154
153
  )
155
154
 
156
- def to_dict(self) -> Dict[str, Any]:
155
+ def to_dict(self) -> dict[str, Any]:
157
156
  """
158
157
  Serializes the component to a dictionary.
159
158
 
@@ -168,7 +167,7 @@ class QueryExpander:
168
167
  )
169
168
 
170
169
  @classmethod
171
- def from_dict(cls, data: Dict[str, Any]) -> "QueryExpander":
170
+ def from_dict(cls, data: dict[str, Any]) -> "QueryExpander":
172
171
  """
173
172
  Deserializes the component from a dictionary.
174
173
 
@@ -181,12 +180,8 @@ class QueryExpander:
181
180
 
182
181
  return default_from_dict(cls, data)
183
182
 
184
- @component.output_types(queries=List[str])
185
- def run(
186
- self,
187
- query: str,
188
- n_expansions: Optional[int] = None,
189
- ) -> Dict[str, List[str]]:
183
+ @component.output_types(queries=list[str])
184
+ def run(self, query: str, n_expansions: Optional[int] = None) -> dict[str, list[str]]:
190
185
  """
191
186
  Expand the input query into multiple semantically similar queries.
192
187
 
@@ -199,11 +194,10 @@ class QueryExpander:
199
194
  If include_original_query=True, the original query will be included in addition
200
195
  to the n_expansions alternative queries.
201
196
  :raises ValueError: If n_expansions is not positive (less than or equal to 0).
202
- :raises RuntimeError: If the component is not warmed up and the chat generator does not support warm up.
203
197
  """
204
198
 
205
- if not self._is_warmed_up and self._supports_warm_up:
206
- raise RuntimeError("The component is not warmed up. Please call the `warm_up` method first.")
199
+ if not self._is_warmed_up:
200
+ self.warm_up()
207
201
 
208
202
  response = {"queries": [query] if self.include_original_query else []}
209
203
 
@@ -252,14 +246,15 @@ class QueryExpander:
252
246
 
253
247
  def warm_up(self):
254
248
  """
255
- Warm up the underlying LLM if it supports it.
249
+ Warm up the LLM provider component.
256
250
  """
257
- if not self._is_warmed_up and self._supports_warm_up:
258
- self.chat_generator.warm_up() # type: ignore[attr-defined]
251
+ if not self._is_warmed_up:
252
+ if hasattr(self.chat_generator, "warm_up"):
253
+ self.chat_generator.warm_up()
259
254
  self._is_warmed_up = True
260
255
 
261
256
  @staticmethod
262
- def _parse_expanded_queries(generator_response: str) -> List[str]:
257
+ def _parse_expanded_queries(generator_response: str) -> list[str]:
263
258
  """
264
259
  Parse the generator response to extract individual expanded queries.
265
260
 
@@ -3,7 +3,7 @@
3
3
  # SPDX-License-Identifier: Apache-2.0
4
4
 
5
5
  from concurrent.futures import ThreadPoolExecutor
6
- from typing import Any, List, Optional
6
+ from typing import Any, Optional
7
7
 
8
8
  from haystack import Document, component, default_from_dict, default_to_dict
9
9
  from haystack.components.embedders.types.protocol import TextEmbedder
@@ -74,13 +74,7 @@ class MultiQueryEmbeddingRetriever:
74
74
  ```
75
75
  """ # noqa E501
76
76
 
77
- def __init__(
78
- self,
79
- *,
80
- retriever: EmbeddingRetriever,
81
- query_embedder: TextEmbedder,
82
- max_workers: int = 3,
83
- ):
77
+ def __init__(self, *, retriever: EmbeddingRetriever, query_embedder: TextEmbedder, max_workers: int = 3) -> None:
84
78
  """
85
79
  Initialize MultiQueryEmbeddingRetriever.
86
80
 
@@ -104,12 +98,8 @@ class MultiQueryEmbeddingRetriever:
104
98
  self.retriever.warm_up()
105
99
  self._is_warmed_up = True
106
100
 
107
- @component.output_types(documents=List[Document])
108
- def run(
109
- self,
110
- queries: List[str],
111
- retriever_kwargs: Optional[dict[str, Any]] = None,
112
- ) -> dict[str, Any]:
101
+ @component.output_types(documents=list[Document])
102
+ def run(self, queries: list[str], retriever_kwargs: Optional[dict[str, Any]] = None) -> dict[str, list[Document]]:
113
103
  """
114
104
  Retrieve documents using multiple queries in parallel.
115
105
 
@@ -123,6 +113,9 @@ class MultiQueryEmbeddingRetriever:
123
113
  seen_contents = set()
124
114
  retriever_kwargs = retriever_kwargs or {}
125
115
 
116
+ if not self._is_warmed_up:
117
+ self.warm_up()
118
+
126
119
  with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
127
120
  queries_results = executor.map(lambda query: self._run_on_thread(query, retriever_kwargs), queries)
128
121
  for result in queries_results:
@@ -137,7 +130,7 @@ class MultiQueryEmbeddingRetriever:
137
130
  docs.sort(key=lambda x: x.score or 0.0, reverse=True)
138
131
  return {"documents": docs}
139
132
 
140
- def _run_on_thread(self, query: str, retriever_kwargs: Optional[dict[str, Any]] = None) -> Optional[List[Document]]:
133
+ def _run_on_thread(self, query: str, retriever_kwargs: Optional[dict[str, Any]] = None) -> Optional[list[Document]]:
141
134
  """
142
135
  Process a single query on a separate thread.
143
136
 
@@ -3,7 +3,7 @@
3
3
  # SPDX-License-Identifier: Apache-2.0
4
4
 
5
5
  from concurrent.futures import ThreadPoolExecutor
6
- from typing import Any, List, Optional
6
+ from typing import Any, Optional
7
7
 
8
8
  from haystack import Document, component, default_from_dict, default_to_dict
9
9
  from haystack.core.serialization import component_to_dict
@@ -57,11 +57,7 @@ class MultiQueryTextRetriever:
57
57
  ```
58
58
  """ # noqa E501
59
59
 
60
- def __init__(
61
- self,
62
- retriever: TextRetriever,
63
- max_workers: int = 3,
64
- ):
60
+ def __init__(self, *, retriever: TextRetriever, max_workers: int = 3) -> None:
65
61
  """
66
62
  Initialize MultiQueryTextRetriever.
67
63
 
@@ -82,11 +78,7 @@ class MultiQueryTextRetriever:
82
78
  self._is_warmed_up = True
83
79
 
84
80
  @component.output_types(documents=list[Document])
85
- def run(
86
- self,
87
- queries: List[str],
88
- retriever_kwargs: Optional[dict[str, Any]] = None,
89
- ) -> dict[str, Any]:
81
+ def run(self, queries: list[str], retriever_kwargs: Optional[dict[str, Any]] = None) -> dict[str, list[Document]]:
90
82
  """
91
83
  Retrieve documents using multiple queries in parallel.
92
84
 
@@ -2,12 +2,12 @@
2
2
  #
3
3
  # SPDX-License-Identifier: Apache-2.0
4
4
 
5
- from copy import deepcopy
6
5
  from dataclasses import replace
7
6
  from datetime import datetime
8
7
  from typing import TYPE_CHECKING, Any, Optional
9
8
 
10
9
  from haystack import logging
10
+ from haystack.core.pipeline.utils import _deepcopy_with_exceptions
11
11
  from haystack.dataclasses.breakpoints import AgentBreakpoint, PipelineSnapshot, PipelineState, ToolBreakpoint
12
12
  from haystack.utils.base_serialization import _serialize_value_with_schema
13
13
  from haystack.utils.misc import _get_output_dir
@@ -44,8 +44,10 @@ def _create_agent_snapshot(
44
44
  """
45
45
  return AgentSnapshot(
46
46
  component_inputs={
47
- "chat_generator": _serialize_value_with_schema(deepcopy(component_inputs["chat_generator"])),
48
- "tool_invoker": _serialize_value_with_schema(deepcopy(component_inputs["tool_invoker"])),
47
+ "chat_generator": _serialize_value_with_schema(
48
+ _deepcopy_with_exceptions(component_inputs["chat_generator"])
49
+ ),
50
+ "tool_invoker": _serialize_value_with_schema(_deepcopy_with_exceptions(component_inputs["tool_invoker"])),
49
51
  },
50
52
  component_visits=component_visits,
51
53
  break_point=agent_breakpoint,