Change structure to structure_factory

griptape-ai · May 3, 2024 · fd616f3 · fd616f3
1 parent 9f0b78e
commit fd616f3
Show file tree

Hide file tree

Showing 6 changed files with 14 additions and 12 deletions.
diff --git a/docs/griptape-framework/drivers/structure-run-drivers.md b/docs/griptape-framework/drivers/structure-run-drivers.md
@@ -32,13 +32,13 @@ joke_coordinator = Pipeline(
     tasks=[
         StructureRunTask(
             driver=LocalStructureRunDriver(
-                structure=joke_teller,
+                structure_factory=lambda: joke_teller,
             ),
         ),
         StructureRunTask(
             "Rewrite this joke: {{ parent_output }}",
             driver=LocalStructureRunDriver(
-                structure=joke_rewriter,
+                structure_factory=lambda: joke_rewriter,
             ),
         ),
     ]

diff --git a/docs/griptape-framework/structures/tasks.md b/docs/griptape-framework/structures/tasks.md
@@ -782,7 +782,7 @@ team = Pipeline(
             Identify key trends, breakthrough technologies, and potential industry impacts.""",
             ),
             driver=LocalStructureRunDriver(
-                structure=researcher,
+                structure_factory=lambda: researcher,
             ),
         ),
         StructureRunTask(
@@ -795,7 +795,7 @@ team = Pipeline(
                 "{{parent_output}}",
             ),
             driver=LocalStructureRunDriver(
-                structure=writer,
+                structure_factory=lambda: writer,
             ),
         ),
     ],

diff --git a/griptape/drivers/structure_run/local_structure_run_driver.py b/griptape/drivers/structure_run/local_structure_run_driver.py
@@ -1,5 +1,5 @@
 from __future__ import annotations
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Callable
 
 from attrs import define, field
 
@@ -12,12 +12,12 @@
 
 @define
 class LocalStructureRunDriver(BaseStructureRunDriver):
-    structure: Structure = field(kw_only=True)
+    structure_factory: Callable[[], Structure] = field(kw_only=True)
 
     def try_run(self, *args: BaseArtifact) -> BaseArtifact:
-        self.structure.run(*[arg.value for arg in args])
+        structure_factory = self.structure_factory().run(*[arg.value for arg in args])
 
-        if self.structure.output_task.output is not None:
-            return self.structure.output_task.output
+        if structure_factory.output_task.output is not None:
+            return structure_factory.output_task.output
         else:
             return InfoArtifact("No output found in response")
diff --git a/tests/unit/drivers/structure_run/test_local_structure_run_driver.py b/tests/unit/drivers/structure_run/test_local_structure_run_driver.py
@@ -10,7 +10,7 @@ class TestLocalStructureRunDriver:
     @pytest.fixture
     def driver(self):
         agent = Agent(prompt_driver=MockPromptDriver(mock_output="agent mock output"))
-        driver = LocalStructureRunDriver(structure=agent)
+        driver = LocalStructureRunDriver(structure_factory=lambda: agent)
 
         return driver
 

diff --git a/tests/unit/tasks/test_structure_run_task.py b/tests/unit/tasks/test_structure_run_task.py
@@ -9,7 +9,7 @@ class TestStructureRunTask:
     def test_run(self):
         agent = Agent(prompt_driver=MockPromptDriver(mock_output="agent mock output"))
         pipeline = Pipeline(prompt_driver=MockPromptDriver(mock_output="pipeline mock output"))
-        driver = LocalStructureRunDriver(structure=agent)
+        driver = LocalStructureRunDriver(structure_factory=lambda: agent)
 
         task = StructureRunTask(driver=driver)
 

diff --git a/tests/unit/tools/test_structure_run_client.py b/tests/unit/tools/test_structure_run_client.py
@@ -11,7 +11,9 @@ def client(self):
         driver = MockPromptDriver()
         agent = Agent(prompt_driver=driver)
 
-        return StructureRunClient(description="foo bar", driver=LocalStructureRunDriver(structure=agent))
+        return StructureRunClient(
+            description="foo bar", driver=LocalStructureRunDriver(structure_factory=lambda: agent)
+        )
 
     def test_run_structure(self, client):
         assert client.run_structure({"values": {"args": "foo bar"}}).value == "mock output"