Coverage for haystack/core/super_component/super

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

3# SPDX-License-Identifier: Apache-2.0

5import functools

6from pathlib import Path

7from types import new_class

8from typing import Any, TypeVar

10from haystack import logging

11from haystack.core.component.component import component

12from haystack.core.pipeline.pipeline import Pipeline

13from haystack.core.pipeline.utils import parse_connect_string

14from haystack.core.serialization import default_from_dict, default_to_dict, generate_qualified_class_name

15from haystack.core.super_component.utils import _delegate_default, _is_compatible

17logger = logging.getLogger(__name__)

19T = TypeVar("T")

22class InvalidMappingTypeError(Exception):

23 """Raised when input or output mappings have invalid types or type conflicts."""

25 pass

28class InvalidMappingValueError(Exception):

29 """Raised when input or output mappings have invalid values or missing components/sockets."""

31 pass

34@component

35class _SuperComponent:

36 def __init__(

37 self,

38 pipeline: Pipeline,

39 input_mapping: dict[str, list[str]] | None = None,

40 output_mapping: dict[str, str] | None = None,

41 ) -> None:

42 """

43 Creates a SuperComponent with optional input and output mappings.

45 :param pipeline: The pipeline instance to be wrapped

46 :param input_mapping: A dictionary mapping component input names to pipeline input socket paths.

47 If not provided, a default input mapping will be created based on all pipeline inputs.

48 Example:

49 ```python

50 input_mapping={

51 "query": ["retriever.query", "prompt_builder.query"],

52 }

53 ```

54 :param output_mapping: A dictionary mapping pipeline output socket paths to component output names.

55 If not provided, a default output mapping will be created based on all pipeline outputs.

56 Example:

57 ```python

58 output_mapping={

59 "retriever.documents": "documents",

60 "generator.replies": "replies",

61 }

62 ```

63 :raises InvalidMappingError: Raised if any mapping is invalid or type conflicts occur

64 :raises ValueError: Raised if no pipeline is provided

65 """

66 if pipeline is None:

67 raise ValueError("Pipeline must be provided to SuperComponent.")

69 self.pipeline: Pipeline = pipeline

71 # Determine input types based on pipeline and mapping

72 pipeline_inputs = self.pipeline.inputs()

73 resolved_input_mapping = (

74 input_mapping if input_mapping is not None else self._create_input_mapping(pipeline_inputs)

75 )

76 self._validate_input_mapping(pipeline_inputs, resolved_input_mapping)

77 input_types = self._resolve_input_types_from_mapping(pipeline_inputs, resolved_input_mapping)

78 # Set input types on the component

79 for input_name, info in input_types.items():

80 component.set_input_type(self, name=input_name, **info)

82 self.input_mapping: dict[str, list[str]] = resolved_input_mapping

83 self._original_input_mapping = input_mapping

85 # Set output types based on pipeline and mapping

86 leaf_pipeline_outputs = self.pipeline.outputs()

87 all_possible_pipeline_outputs = self.pipeline.outputs(include_components_with_connected_outputs=True)

89 resolved_output_mapping = (

90 output_mapping if output_mapping is not None else self._create_output_mapping(leaf_pipeline_outputs)

91 )

92 self._validate_output_mapping(all_possible_pipeline_outputs, resolved_output_mapping)

93 output_types = self._resolve_output_types_from_mapping(all_possible_pipeline_outputs, resolved_output_mapping)

94 # Set output types on the component

95 component.set_output_types(self, **output_types)

96 self.output_mapping: dict[str, str] = resolved_output_mapping

97 self._original_output_mapping = output_mapping

99 def warm_up(self) -> None:

100 """

101 Warms up the SuperComponent by warming up the wrapped pipeline.

102 """

103 self.pipeline.warm_up()

104

105 async def warm_up_async(self) -> None:

106 """

107 Warms up the SuperComponent by warming up the wrapped pipeline on the serving event loop.

108 """

109 await self.pipeline.warm_up_async()

110

111 def close(self) -> None:

112 """

113 Releases the synchronous resources held by the wrapped pipeline's components.

114 """

115 self.pipeline.close()

116

117 async def close_async(self) -> None:

118 """

119 Releases the async resources held by the wrapped pipeline's components.

120 """

121 await self.pipeline.close_async()

122

123 def run(self, **kwargs: Any) -> dict[str, Any]:

124 """

125 Runs the wrapped pipeline with the provided inputs.

126

127 Steps:

128 1. Maps the inputs from kwargs to pipeline component inputs

129 2. Runs the pipeline

130 3. Maps the pipeline outputs to the SuperComponent's outputs

131

132 :param kwargs: Keyword arguments matching the SuperComponent's input names

133 :returns:

134 Dictionary containing the SuperComponent's output values

135 """

136 # `is not`, not `!=`: numpy/pandas/torch override `__ne__` element-wise and would crash here.

137 filtered_inputs = {param: value for param, value in kwargs.items() if value is not _delegate_default}

138 pipeline_inputs = self._map_explicit_inputs(input_mapping=self.input_mapping, inputs=filtered_inputs)

139 include_outputs_from = self._get_include_outputs_from()

140 pipeline_outputs = self.pipeline.run(data=pipeline_inputs, include_outputs_from=include_outputs_from)

141 return self._map_explicit_outputs(pipeline_outputs, self.output_mapping)

142

143 def _get_include_outputs_from(self) -> set[str]:

144 # Collecting the component names from output_mapping

145 return {self._split_component_path(path)[0] for path in self.output_mapping.keys()}

146

147 async def run_async(self, **kwargs: Any) -> dict[str, Any]:

148 """

149 Runs the wrapped pipeline with the provided inputs async.

150

151 Steps:

152 1. Maps the inputs from kwargs to pipeline component inputs

153 2. Runs the pipeline async

154 3. Maps the pipeline outputs to the SuperComponent's outputs

155

156 :param kwargs: Keyword arguments matching the SuperComponent's input names

157 :returns:

158 Dictionary containing the SuperComponent's output values

159 """

160 # `is not`, not `!=`: numpy/pandas/torch override `__ne__` element-wise and would crash here.

161 filtered_inputs = {param: value for param, value in kwargs.items() if value is not _delegate_default}

162 pipeline_inputs = self._map_explicit_inputs(input_mapping=self.input_mapping, inputs=filtered_inputs)

163 include_outputs_from = self._get_include_outputs_from()

164 pipeline_outputs = await self.pipeline.run_async(

165 data=pipeline_inputs, include_outputs_from=include_outputs_from

166 )

167 return self._map_explicit_outputs(pipeline_outputs, self.output_mapping)

168

169 @staticmethod

170 def _split_component_path(path: str) -> tuple[str, str]:

171 """

172 Splits a component path into a component name and a socket name.

173

174 :param path: A string in the format "component_name.socket_name".

175 :returns:

176 A tuple containing (component_name, socket_name).

177 :raises InvalidMappingValueError:

178 If the path format is incorrect.

179 """

180 comp_name, socket_name = parse_connect_string(path)

181 if socket_name is None:

182 raise InvalidMappingValueError(f"Invalid path format: '{path}'. Expected 'component_name.socket_name'.")

183 return comp_name, socket_name

184

185 def _validate_input_mapping(

186 self, pipeline_inputs: dict[str, dict[str, Any]], input_mapping: dict[str, list[str]]

187 ) -> None:

188 """

189 Validates the input mapping to ensure that specified components and sockets exist in the pipeline.

190

191 :param pipeline_inputs: A dictionary containing pipeline input specifications.

192 :param input_mapping: A dictionary mapping wrapper input names to pipeline socket paths.

193 :raises InvalidMappingTypeError:

194 If the input mapping is of invalid type or contains invalid types.

195 :raises InvalidMappingValueError:

196 If the input mapping contains nonexistent components or sockets.

197 """

198 if not isinstance(input_mapping, dict):

199 raise InvalidMappingTypeError("input_mapping must be a dictionary")

200

201 for wrapper_input_name, pipeline_input_paths in input_mapping.items():

202 if not isinstance(pipeline_input_paths, list):

203 raise InvalidMappingTypeError(f"Input paths for '{wrapper_input_name}' must be a list of strings.")

204 for path in pipeline_input_paths:

205 comp_name, socket_name = self._split_component_path(path)

206 if comp_name not in pipeline_inputs:

207 raise InvalidMappingValueError(

208 f"Component '{comp_name}' not found in pipeline inputs.\n"

209 f"Available components: {list(pipeline_inputs.keys())}"

210 )

211 if socket_name not in pipeline_inputs[comp_name]:

212 raise InvalidMappingValueError(

213 f"Input socket '{socket_name}' not found in component '{comp_name}'.\n"

214 f"Available inputs for '{comp_name}': {list(pipeline_inputs[comp_name].keys())}"

215 )

216

217 def _resolve_input_types_from_mapping(

218 self, pipeline_inputs: dict[str, dict[str, Any]], input_mapping: dict[str, list[str]]

219 ) -> dict[str, dict[str, Any]]:

220 """

221 Resolves and validates input types based on the provided input mapping.

222

223 This function ensures that all mapped pipeline inputs are compatible, consolidating types

224 when multiple mappings exist. It also determines whether an input is mandatory or has a default value.

225

226 :param pipeline_inputs: A dictionary containing pipeline input specifications.

227 :param input_mapping: A dictionary mapping SuperComponent inputs to pipeline socket paths.

228 :returns:

229 A dictionary specifying the resolved input types and their properties.

230 :raises InvalidMappingTypeError:

231 If the input mapping contains incompatible types.

232 """

233 aggregated_inputs: dict[str, dict[str, Any]] = {}

234 for wrapper_input_name, pipeline_input_paths in input_mapping.items():

235 for path in pipeline_input_paths:

236 comp_name, socket_name = self._split_component_path(path)

237 socket_info = pipeline_inputs[comp_name][socket_name]

238

239 # Add to aggregated inputs

240 existing_socket_info = aggregated_inputs.get(wrapper_input_name)

241 if existing_socket_info is None:

242 aggregated_inputs[wrapper_input_name] = {"type": socket_info["type"]}

243 if not socket_info["is_mandatory"]:

244 aggregated_inputs[wrapper_input_name]["default"] = _delegate_default

245 continue

246

247 is_compatible, common_type = _is_compatible(existing_socket_info["type"], socket_info["type"])

248

249 if not is_compatible:

250 raise InvalidMappingTypeError(

251 f"Type conflict for input '{socket_name}' from component '{comp_name}'. "

252 f"Existing type: {existing_socket_info['type']}, new type: {socket_info['type']}."

253 )

254

255 # Use the common type for the aggregated input

256 aggregated_inputs[wrapper_input_name]["type"] = common_type

257

258 # If any socket requires mandatory inputs then the aggregated input is also considered mandatory.

259 # So we use the type of the mandatory input and remove the default value if it exists.

260 if socket_info["is_mandatory"]:

261 aggregated_inputs[wrapper_input_name].pop("default", None)

262

263 return aggregated_inputs

264

265 @staticmethod

266 def _create_input_mapping(pipeline_inputs: dict[str, dict[str, Any]]) -> dict[str, list[str]]:

267 """

268 Create an input mapping from pipeline inputs.

269

270 :param pipeline_inputs: Dictionary of pipeline input specifications

271 :returns:

272 Dictionary mapping SuperComponent input names to pipeline socket paths

273 """

274 input_mapping: dict[str, list[str]] = {}

275 for comp_name, inputs_dict in pipeline_inputs.items():

276 for socket_name in inputs_dict.keys():

277 existing_socket_info = input_mapping.get(socket_name)

278 if existing_socket_info is None:

279 input_mapping[socket_name] = [f"{comp_name}.{socket_name}"]

280 continue

281 input_mapping[socket_name].append(f"{comp_name}.{socket_name}")

282 return input_mapping

283

284 def _validate_output_mapping(

285 self, pipeline_outputs: dict[str, dict[str, Any]], output_mapping: dict[str, str]

286 ) -> None:

287 """

288 Validates the output mapping to ensure that specified components and sockets exist in the pipeline.

289

290 :param pipeline_outputs: A dictionary containing pipeline output specifications.

291 :param output_mapping: A dictionary mapping pipeline socket paths to wrapper output names.

292 :raises InvalidMappingTypeError:

293 If the output mapping is of invalid type or contains invalid types.

294 :raises InvalidMappingValueError:

295 If the output mapping contains nonexistent components or sockets.

296 """

297 for pipeline_output_path, wrapper_output_name in output_mapping.items():

298 if not isinstance(wrapper_output_name, str):

299 raise InvalidMappingTypeError("Output names in output_mapping must be strings.")

300 comp_name, socket_name = self._split_component_path(pipeline_output_path)

301 if comp_name not in pipeline_outputs:

302 raise InvalidMappingValueError(f"Component '{comp_name}' not found among pipeline outputs.")

303 if socket_name not in pipeline_outputs[comp_name]:

304 raise InvalidMappingValueError(f"Output socket '{socket_name}' not found in component '{comp_name}'.")

305

306 def _resolve_output_types_from_mapping(

307 self, pipeline_outputs: dict[str, dict[str, Any]], output_mapping: dict[str, str]

308 ) -> dict[str, Any]:

309 """

310 Resolves and validates output types based on the provided output mapping.

311

312 This function ensures that all mapped pipeline outputs are correctly assigned to

313 the corresponding SuperComponent outputs while preventing duplicate output names.

314

315 :param pipeline_outputs: A dictionary containing pipeline output specifications.

316 :param output_mapping: A dictionary mapping pipeline output socket paths to SuperComponent output names.

317 :returns:

318 A dictionary mapping SuperComponent output names to their resolved types.

319 :raises InvalidMappingValueError:

320 If the output mapping contains duplicate output names.

321 """

322 resolved_outputs = {}

323 for pipeline_output_path, wrapper_output_name in output_mapping.items():

324 comp_name, socket_name = self._split_component_path(pipeline_output_path)

325 if wrapper_output_name in resolved_outputs:

326 raise InvalidMappingValueError(f"Duplicate output name '{wrapper_output_name}' in output_mapping.")

327 resolved_outputs[wrapper_output_name] = pipeline_outputs[comp_name][socket_name]["type"]

328 return resolved_outputs

329

330 @staticmethod

331 def _create_output_mapping(pipeline_outputs: dict[str, dict[str, Any]]) -> dict[str, str]:

332 """

333 Create an output mapping from pipeline outputs.

334

335 :param pipeline_outputs: Dictionary of pipeline output specifications

336 :returns:

337 Dictionary mapping pipeline socket paths to SuperComponent output names

338 :raises InvalidMappingValueError:

339 If there are output name conflicts between components

340 """

341 output_mapping = {}

342 used_output_names: set[str] = set()

343 for comp_name, outputs_dict in pipeline_outputs.items():

344 for socket_name in outputs_dict.keys():

345 if socket_name in used_output_names:

346 raise InvalidMappingValueError(

347 f"Output name conflict: '{socket_name}' is produced by multiple components. "

348 "Please provide an output_mapping to resolve this conflict."

349 )

350 used_output_names.add(socket_name)

351 output_mapping[f"{comp_name}.{socket_name}"] = socket_name

352 return output_mapping

353

354 def _map_explicit_inputs(

355 self, input_mapping: dict[str, list[str]], inputs: dict[str, Any]

356 ) -> dict[str, dict[str, Any]]:

357 """

358 Map inputs according to explicit input mapping.

359

360 :param input_mapping: Mapping configuration for inputs

361 :param inputs: Input arguments provided to wrapper

362 :return: Dictionary of mapped pipeline inputs

363 """

364 pipeline_inputs: dict[str, dict[str, Any]] = {}

365 for wrapper_input_name, pipeline_input_paths in input_mapping.items():

366 if wrapper_input_name not in inputs:

367 continue

368

369 for socket_path in pipeline_input_paths:

370 comp_name, input_name = self._split_component_path(socket_path)

371 if comp_name not in pipeline_inputs:

372 pipeline_inputs[comp_name] = {}

373 pipeline_inputs[comp_name][input_name] = inputs[wrapper_input_name]

374

375 return pipeline_inputs

376

377 def _map_explicit_outputs(

378 self, pipeline_outputs: dict[str, dict[str, Any]], output_mapping: dict[str, str]

379 ) -> dict[str, Any]:

380 """

381 Map outputs according to explicit output mapping.

382

383 :param pipeline_outputs: Raw outputs from pipeline execution

384 :param output_mapping: Output mapping configuration

385 :return: Dictionary of mapped outputs

386 """

387 outputs: dict[str, Any] = {}

388 for pipeline_output_path, wrapper_output_name in output_mapping.items():

389 comp_name, socket_name = self._split_component_path(pipeline_output_path)

390 if comp_name in pipeline_outputs and socket_name in pipeline_outputs[comp_name]:

391 outputs[wrapper_output_name] = pipeline_outputs[comp_name][socket_name]

392 return outputs

393

394 def _to_super_component_dict(self) -> dict[str, Any]:

395 """

396 Convert to a SuperComponent dictionary representation.

397

398 :return: Dictionary containing serialized SuperComponent data

399 """

400 serialized_pipeline = self.pipeline.to_dict()

401 serialized = default_to_dict(

402 self,

403 pipeline=serialized_pipeline,

404 input_mapping=self._original_input_mapping,

405 output_mapping=self._original_output_mapping,

406 )

407 serialized["type"] = generate_qualified_class_name(SuperComponent)

408 return serialized

409

410

411@component

412class SuperComponent(_SuperComponent):

413 """

414 A class for creating super components that wrap around a Pipeline.

415

416 This component allows for remapping of input and output socket names between the wrapped pipeline and the

417 SuperComponent's input and output names. This is useful for creating higher-level components that abstract

418 away the details of the wrapped pipeline.

419

420 ### Usage example

421

422 ```python

423 from haystack import Pipeline, SuperComponent

424 from haystack.components.generators.chat import OpenAIResponsesChatGenerator

425 from haystack.components.builders import ChatPromptBuilder

426 from haystack.components.retrievers import InMemoryBM25Retriever

427 from haystack.dataclasses.chat_message import ChatMessage

428 from haystack.document_stores.in_memory import InMemoryDocumentStore

429 from haystack.dataclasses import Document

430

431 document_store = InMemoryDocumentStore()

432 documents = [

433 Document(content="Paris is the capital of France."),

434 Document(content="London is the capital of England."),

435 ]

436 document_store.write_documents(documents)

437

438 prompt_template = [

439 ChatMessage.from_user(

440 '''

441 According to the following documents:

442 {% for document in documents %}

443 {{document.content}}

444 {% endfor %}

445 Answer the given question: {{query}}

446 Answer:

447 '''

448 )

449 ]

450

451 prompt_builder = ChatPromptBuilder(template=prompt_template, required_variables="*")

452

453 pipeline = Pipeline()

454 pipeline.add_component("retriever", InMemoryBM25Retriever(document_store=document_store))

455 pipeline.add_component("prompt_builder", prompt_builder)

456 pipeline.add_component("llm", OpenAIResponsesChatGenerator())

457 pipeline.connect("retriever.documents", "prompt_builder.documents")

458 pipeline.connect("prompt_builder.prompt", "llm.messages")

459

460 # Create a super component with simplified input/output mapping

461 wrapper = SuperComponent(

462 pipeline=pipeline,

463 input_mapping={

464 "query": ["retriever.query", "prompt_builder.query"],

465 },

466 output_mapping={"llm.replies": "replies"}

467 )

468

469 # Run the pipeline with simplified interface

470 result = wrapper.run(query="What is the capital of France?")

471 print(result)

472 {'replies': [ChatMessage(_role=<ChatRole.ASSISTANT: 'assistant'>,

473 _content=[TextContent(text='The capital of France is Paris.')],...)

474 ```

475

476 """

477

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

479 """

480 Serializes the SuperComponent into a dictionary.

481

482 :returns:

483 Dictionary with serialized data.

484 """

485 return self._to_super_component_dict()

486

487 @classmethod

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

489 """

490 Deserializes the SuperComponent from a dictionary.

491

492 :param data: The dictionary to deserialize from.

493 :returns:

494 The deserialized SuperComponent.

495 """

496 # `is_pipeline_async` is a legacy key kept only for backward compatibility.

497 data["init_parameters"].pop("is_pipeline_async", None)

498 pipeline = Pipeline.from_dict(data["init_parameters"]["pipeline"])

499 data["init_parameters"]["pipeline"] = pipeline

500 return default_from_dict(cls, data)

501

502 def show(self, server_url: str = "https://mermaid.ink", params: dict | None = None, timeout: int = 30) -> None:

503 """

504 Display an image representing this SuperComponent's underlying pipeline in a Jupyter notebook.

505

506 This function generates a diagram of the Pipeline using a Mermaid server and displays it directly in

507 the notebook.

508

509 :param server_url:

510 The base URL of the Mermaid server used for rendering (default: 'https://mermaid.ink').

511 See https://github.com/jihchi/mermaid.ink and https://github.com/mermaid-js/mermaid-live-editor for more

512 info on how to set up your own Mermaid server.

513

514 :param params:

515 Dictionary of customization parameters to modify the output. Refer to Mermaid documentation for more details

516 Supported keys:

517 - format: Output format ('img', 'svg', or 'pdf'). Default: 'img'.

518 - type: Image type for /img endpoint ('jpeg', 'png', 'webp'). Default: 'png'.

519 - theme: Mermaid theme ('default', 'neutral', 'dark', 'forest'). Default: 'neutral'.

520 - bgColor: Background color in hexadecimal (e.g., 'FFFFFF') or named format (e.g., '!white').

521 - width: Width of the output image (integer).

522 - height: Height of the output image (integer).

523 - scale: Scaling factor (1–3). Only applicable if 'width' or 'height' is specified.

524 - fit: Whether to fit the diagram size to the page (PDF only, boolean).

525 - paper: Paper size for PDFs (e.g., 'a4', 'a3'). Ignored if 'fit' is true.

526 - landscape: Landscape orientation for PDFs (boolean). Ignored if 'fit' is true.

527

528 :param timeout:

529 Timeout in seconds for the request to the Mermaid server.

530

531 :raises PipelineDrawingError:

532 If the function is called outside of a Jupyter notebook or if there is an issue with rendering.

533 """

534 self.pipeline.show(server_url=server_url, params=params, timeout=timeout)

535

536 def draw(

537 self, path: Path, server_url: str = "https://mermaid.ink", params: dict | None = None, timeout: int = 30

538 ) -> None:

539 """

540 Save an image representing this SuperComponent's underlying pipeline to the specified file path.

541

542 This function generates a diagram of the Pipeline using the Mermaid server and saves it to the provided path.

543

544 :param path:

545 The file path where the generated image will be saved.

546 :param server_url:

547 The base URL of the Mermaid server used for rendering (default: 'https://mermaid.ink').

548 See https://github.com/jihchi/mermaid.ink and https://github.com/mermaid-js/mermaid-live-editor for more

549 info on how to set up your own Mermaid server.

550 :param params:

551 Dictionary of customization parameters to modify the output. Refer to Mermaid documentation for more details

552 Supported keys:

553 - format: Output format ('img', 'svg', or 'pdf'). Default: 'img'.

554 - type: Image type for /img endpoint ('jpeg', 'png', 'webp'). Default: 'png'.

555 - theme: Mermaid theme ('default', 'neutral', 'dark', 'forest'). Default: 'neutral'.

556 - bgColor: Background color in hexadecimal (e.g., 'FFFFFF') or named format (e.g., '!white').

557 - width: Width of the output image (integer).

558 - height: Height of the output image (integer).

559 - scale: Scaling factor (1–3). Only applicable if 'width' or 'height' is specified.

560 - fit: Whether to fit the diagram size to the page (PDF only, boolean).

561 - paper: Paper size for PDFs (e.g., 'a4', 'a3'). Ignored if 'fit' is true.

562 - landscape: Landscape orientation for PDFs (boolean). Ignored if 'fit' is true.

563

564 :param timeout:

565 Timeout in seconds for the request to the Mermaid server.

566

567 :raises PipelineDrawingError:

568 If there is an issue with rendering or saving the image.

569 """

570 self.pipeline.draw(path=path, server_url=server_url, params=params, timeout=timeout)

571

572

573def super_component(cls: type[T]) -> type[T]:

574 """

575 Decorator that converts a class into a SuperComponent.

576

577 This decorator:

578 1. Creates a new class that inherits from SuperComponent

579 2. Copies all methods and attributes from the original class

580 3. Adds initialization logic to properly set up the SuperComponent

581

582 The decorated class should define:

583 - pipeline: A Pipeline instance in the __init__ method

584 - input_mapping: Dictionary mapping component inputs to pipeline inputs (optional)

585 - output_mapping: Dictionary mapping pipeline outputs to component outputs (optional)

586 """

587 logger.debug("Registering {cls} as a super_component", cls=cls)

588

589 # Store the original __init__ method

590 original_init = cls.__init__

591

592 # Create a new __init__ method that will initialize both the original class and SuperComponent

593 def init_wrapper(self: Any, *args: Any, **kwargs: Any) -> None:

594 # Call the original __init__ to set up pipeline and mappings

595 original_init(self, *args, **kwargs)

596

597 # Verify required attributes

598 if not hasattr(self, "pipeline"):

599 raise ValueError(f"Class {cls.__name__} decorated with @super_component must define a 'pipeline' attribute")

600

601 # Initialize SuperComponent

602 _SuperComponent.__init__(

603 self,

604 pipeline=self.pipeline,

605 input_mapping=getattr(self, "input_mapping", None),

606 output_mapping=getattr(self, "output_mapping", None),

607 )

608

609 # Preserve original init's signature for IDEs/docs/tools

610 init_wrapper = functools.wraps(original_init)(init_wrapper)

611

612 # Function to copy namespace from the original class

613 def copy_class_namespace(namespace: dict[str, Any]) -> None:

614 """Copy all attributes from the original class except special ones."""

615 for key, val in dict(cls.__dict__).items():

616 # Skip special attributes that should be recreated

617 if key in ("__dict__", "__weakref__"):

618 continue

619

620 # Override __init__ with our wrapper

621 if key == "__init__":

622 namespace["__init__"] = init_wrapper

623 continue

624

625 namespace[key] = val

626

627 # Create a new class inheriting from SuperComponent with the original methods

628 # We use (SuperComponent,) + cls.__bases__ to make the new class inherit from

629 # SuperComponent and all the original class's bases

630 new_cls = new_class(cls.__name__, (_SuperComponent,) + cls.__bases__, {}, copy_class_namespace)

631

632 # Copy other class attributes

633 new_cls.__module__ = cls.__module__

634 new_cls.__qualname__ = cls.__qualname__

635 new_cls.__doc__ = cls.__doc__

636

637 # Apply the component decorator to the new class

638 return component(new_cls)

Coverage for haystack/core/super_component/super_component.py: 96%

195 statements