Change StructureRunTask parent to PromptTask

griptape-ai · Oct 7, 2024 · bbbfbc8 · bbbfbc8
1 parent 2f5c871
commit bbbfbc8
Show file tree

Hide file tree

Showing 6 changed files with 8 additions and 131 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -63,6 +63,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - **BREAKING**: `BaseExtractionEngine` no longer catches exceptions and returns `ErrorArtifact`s.
 - **BREAKING**: `JsonExtractionEngine.template_schema` is now required.
 - **BREAKING**: `CsvExtractionEngine.column_names` is now required.
+- `StructureRunTask` now inherits from `PromptTask`.
 - `JsonExtractionEngine.extract_artifacts` now returns a `ListArtifact[JsonArtifact]`.
 - `CsvExtractionEngine.extract_artifacts` now returns a `ListArtifact[CsvRowArtifact]`.
 - Remove `manifest.yml` requirements for custom tool creation.

diff --git a/griptape/tasks/__init__.py b/griptape/tasks/__init__.py
@@ -1,6 +1,5 @@
 from .base_task import BaseTask
 from .base_text_input_task import BaseTextInputTask
-from .base_multi_text_input_task import BaseMultiTextInputTask
 from .prompt_task import PromptTask
 from .actions_subtask import ActionsSubtask
 from .toolkit_task import ToolkitTask
@@ -23,7 +22,6 @@
 __all__ = [
     "BaseTask",
     "BaseTextInputTask",
-    "BaseMultiTextInputTask",
     "PromptTask",
     "ActionsSubtask",
     "ToolkitTask",

diff --git a/griptape/tasks/base_multi_text_input_task.py b/griptape/tasks/base_multi_text_input_task.py
diff --git a/griptape/tasks/structure_run_task.py b/griptape/tasks/structure_run_task.py
@@ -4,15 +4,16 @@
 
 from attrs import define, field
 
-from griptape.tasks import BaseMultiTextInputTask
+from griptape.artifacts.list_artifact import ListArtifact
+from griptape.tasks.prompt_task import PromptTask
 
 if TYPE_CHECKING:
     from griptape.artifacts import BaseArtifact
     from griptape.drivers.structure_run.base_structure_run_driver import BaseStructureRunDriver
 
 
 @define
-class StructureRunTask(BaseMultiTextInputTask):
+class StructureRunTask(PromptTask):
     """Task to run a Structure.
 
     Attributes:
@@ -22,4 +23,7 @@ class StructureRunTask(BaseMultiTextInputTask):
     driver: BaseStructureRunDriver = field(kw_only=True)
 
     def run(self) -> BaseArtifact:
-        return self.driver.run(*self.input)
+        if isinstance(self.input, ListArtifact):
+            return self.driver.run(*self.input.value)
+        else:
+            return self.driver.run(self.input)
diff --git a/tests/mocks/mock_multi_text_input_task.py b/tests/mocks/mock_multi_text_input_task.py
diff --git a/tests/unit/tasks/test_base_multi_text_input_task.py b/tests/unit/tasks/test_base_multi_text_input_task.py