More fixes

griptape-ai · May 2, 2024 · 68fce49 · 68fce49
1 parent ed82af3
commit 68fce49
Show file tree

Hide file tree

Showing 6 changed files with 35 additions and 15 deletions.
diff --git a/docs/griptape-framework/drivers/structure-run-drivers.md b/docs/griptape-framework/drivers/structure-run-drivers.md
@@ -4,6 +4,8 @@ When combined with the [StructureRunTask](../../reference/griptape/tasks/structu
 
 ## Local Structure Run Driver
 
+The [LocalStructureRunDriver](../../reference/griptape/drivers/structure-run/local-structure-run-driver.md) is used to run Griptape Structures in the same runtime environment as the code that is running the Structure.
+
 ```python
 from griptape.drivers import LocalStructureRunDriver
 from griptape.rules import Rule
@@ -47,11 +49,15 @@ joke_coordinator.run("Tell me a joke")
 
 ## Griptape Cloud Structure Run Driver
 
+The [GriptapeCloudStructureRunDriver](../../reference/griptape/drivers/structure-run/griptape-cloud-structure-run-driver.md)  is used to run Griptape Structures in the Griptape Cloud.
+
+
 ```python
 import os
 
 from griptape.drivers import GriptapeCloudStructureRunDriver, LocalStructureRunDriver
 from griptape.structures import Pipeline, Agent
+from griptape.rules import Rule
 from griptape.tasks import StructureRunTask
 
 base_url = os.environ["GRIPTAPE_CLOUD_BASE_URL"]
@@ -62,8 +68,19 @@ structure_id = os.environ["GRIPTAPE_CLOUD_STRUCTURE_ID"]
 pipeline = Pipeline(
     tasks=[
         StructureRunTask(
-            "Think of a question related to RAG.",
-            driver=LocalStructureRunDriver(structure=Agent()),
+            "Think of a question related to Retrieval Augmented Generation.",
+            driver=LocalStructureRunDriver(
+                structure=Agent(
+                    rules=[
+                        Rule(
+                            value="You are an expert in Retrieval Augmented Generation.",
+                        ),
+                        Rule(
+                            value="Only output your answer, no other information.",
+                        ),
+                    ]
+                )
+            ),
         ),
         StructureRunTask(
             "{{ parent_output }}",

diff --git a/griptape/drivers/structure_run/base_structure_run_driver.py b/griptape/drivers/structure_run/base_structure_run_driver.py
@@ -2,16 +2,13 @@
 
 from attrs import define
 
-from griptape.artifacts import BaseArtifact, ErrorArtifact
+from griptape.artifacts import BaseArtifact
 
 
 @define
 class BaseStructureRunDriver(ABC):
     def run(self, *args: BaseArtifact) -> BaseArtifact:
-        try:
-            return self.try_run(*args)
-        except Exception as e:
-            return ErrorArtifact(str(e))
+        return self.try_run(*args)
 
     @abstractmethod
     def try_run(self, *args: BaseArtifact) -> BaseArtifact:

diff --git a/griptape/drivers/structure_run/local_structure_run_driver.py b/griptape/drivers/structure_run/local_structure_run_driver.py
@@ -14,8 +14,8 @@
 class LocalStructureRunDriver(BaseStructureRunDriver):
     structure: Structure = field(kw_only=True)
 
-    def try_run(self, *args) -> BaseArtifact:
-        self.structure.run(*args)
+    def try_run(self, *args: BaseArtifact) -> BaseArtifact:
+        self.structure.run(*[arg.value for arg in args])
 
         if self.structure.output_task.output is not None:
             return self.structure.output_task.output

diff --git a/griptape/tools/structure_run_client/tool.py b/griptape/tools/structure_run_client/tool.py
@@ -29,6 +29,6 @@ class StructureRunClient(BaseTool):
         }
     )
     def run_structure(self, params: dict) -> BaseArtifact:
-        args: str = params["values"]["args"]
+        args: list[str] = params["values"]["args"]
 
         return self.driver.run(*[TextArtifact(arg) for arg in args])
diff --git a/tests/unit/drivers/structure_run/test_griptape_cloud_structure_run_driver.py b/tests/unit/drivers/structure_run/test_griptape_cloud_structure_run_driver.py
@@ -12,10 +12,16 @@ def driver(self, mocker):
         mocker.patch("requests.post", return_value=mock_response)
 
         mock_response = mocker.Mock()
-        mock_response.json.return_value = {"description": "fizz buzz", "output": "fooey booey", "status": "SUCCEEDED"}
+        mock_response.json.return_value = {
+            "description": "fizz buzz",
+            "output": TextArtifact("foo bar").to_dict(),
+            "status": "SUCCEEDED",
+        }
         mocker.patch("requests.get", return_value=mock_response)
 
         return GriptapeCloudStructureRunDriver(base_url="https://api.griptape.ai", api_key="foo bar", structure_id="1")
 
     def test_run(self, driver):
-        assert isinstance(driver.run("foo bar"), TextArtifact)
+        result = driver.run("foo bar")
+        assert isinstance(result, TextArtifact)
+        assert result.value == "foo bar"
diff --git a/...ols/test_griptape_structure_run_client.py → ...s/unit/tools/test_structure_run_client.py b/...ols/test_griptape_structure_run_client.py → ...s/unit/tools/test_structure_run_client.py
@@ -1,17 +1,17 @@
 import pytest
 from griptape.drivers.structure_run.local_structure_run_driver import LocalStructureRunDriver
-from griptape.tools import GriptapeStructureRunClient
+from griptape.tools import StructureRunClient
 from griptape.structures import Agent
 from tests.mocks.mock_prompt_driver import MockPromptDriver
 
 
-class TestGriptapeStructureRunClient:
+class TestStructureRunClient:
     @pytest.fixture
     def client(self):
         driver = MockPromptDriver()
         agent = Agent(prompt_driver=driver)
 
-        return GriptapeStructureRunClient(description="foo bar", driver=LocalStructureRunDriver(structure=agent))
+        return StructureRunClient(description="foo bar", driver=LocalStructureRunDriver(structure=agent))
 
     def test_run_structure(self, client):
         assert client.run_structure({"values": {"args": "foo bar"}}).value == "mock output"