fix: Jinja2-based ComponentTools in pipelines properly handle deepcopy (#234)

Author: mathislucka

haystack_experimental/tools/component_tool.py

@@ -2,6 +2,7 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+from copy import copy, deepcopy
 from dataclasses import fields, is_dataclass
 from inspect import getdoc
 from typing import Any, Callable, Dict, Optional, Union, get_args, get_origin
@@ -392,3 +393,29 @@ def _create_property_schema(self, python_type: Any, description: str, default: A
             schema["default"] = default
 
         return schema
+
+    def __deepcopy__(self, memo: Dict[Any, Any]) -> "ComponentTool":
+        # Jinja2 templates throw an Exception when we deepcopy them (see https://github.com/pallets/jinja/issues/758)
+        # When we use a ComponentTool in a pipeline at runtime, we deepcopy the tool
+        # We overwrite ComponentTool.__deepcopy__ to fix this in experimental until a more comprehensive fix is merged.
+        # We track the issue here: https://github.com/deepset-ai/haystack/issues/9011
+        result = copy(self)
+
+        # Add the object to the memo dictionary to handle circular references
+        memo[id(self)] = result
+
+        # Deep copy all attributes with exception handling
+        for key, value in self.__dict__.items():
+            try:
+                # Try to deep copy the attribute
+                setattr(result, key, deepcopy(value, memo))
+            except TypeError:
+                # Fall back to using the original attribute for components that use Jinja2-templates
+                logger.debug(
+                    "deepcopy of ComponentTool {tool_name} failed. Using original attribute '{attribute}' instead.",
+                    tool_name=self.name,
+                    attribute=key,
+                )
+                setattr(result, key, getattr(self, key))
+
+        return result

test/tools/test_component_tool.py

@@ -4,12 +4,14 @@
 
 import json
 import os
+from copy import deepcopy
 from dataclasses import dataclass
 from typing import Dict, List
 
 import pytest
 
 from haystack import Pipeline, component
+from haystack.components.builders import PromptBuilder
 from haystack.components.generators.chat import OpenAIChatGenerator
 from haystack.components.websearch.serper_dev import SerperDevWebSearch
 from haystack.dataclasses import ChatMessage, ChatRole, Document
@@ -609,3 +611,21 @@ def test_pipeline_component_fails(self):
         # thus can't be used as tool
         with pytest.raises(ValueError, match="Component has been added to a pipeline"):
             ComponentTool(component=component)
+
+    def test_deepcopy_with_jinja_based_component(self):
+        # Jinja2 templates throw an Exception when we deepcopy them (see https://github.com/pallets/jinja/issues/758)
+        # When we use a ComponentTool in a pipeline at runtime, we deepcopy the tool
+        # We overwrite ComponentTool.__deepcopy__ to fix this in experimental until a more comprehensive fix is merged.
+        # We track the issue here: https://github.com/deepset-ai/haystack/issues/9011
+
+        builder = PromptBuilder("{{query}}")
+
+        tool = ComponentTool(component=builder)
+        result = tool.function(query="Hello")
+
+        tool_copy = deepcopy(tool)
+
+        result_from_copy = tool_copy.function(query="Hello")
+
+        assert "prompt" in result_from_copy
+        assert result_from_copy["prompt"] == result["prompt"]