Repository: griptape-ai/griptape

Issue:
Can't Set Allowlist on DateTimeTool
https://github.com/griptape-ai/griptape/pull/1611 introduced some new activities for the DateTimeTool and `denylist`'d the old one. That breaks an example like this:

```python
from griptape.structures import Agent
from griptape.tools import DateTimeTool

ollama_agent = Agent(
    tools=[
        DateTimeTool(allowlist=["get_current_datetime"]),
    ],
)

ollama_agent.run("What is today's date?")
```

```
Traceback (most recent call last):
  File "/Users/collindutter/Documents/griptape/griptape-playground/src/misc/datetime_mistral.py", line 6, in <module>
    DateTimeTool(allowlist=["get_current_datetime"]),
    ~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<attrs generated init griptape.tools.date_time.tool.DateTimeTool>", line 30, in __init__
    __attr_validator_allowlist(self, __attr_allowlist, self.allowlist)
    ~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/collindutter/Documents/griptape/griptape/griptape/mixins/activity_mixin.py", line 35, in validate_allowlist
    raise ValueError("can't have both allowlist and denylist specified")
ValueError: can't have both allowlist and denylist specified
```

Maybe we can remove the restriction on preventing allowlist and denylist? Go for a denied-by-default model?

Discussion:

Solution patch:
diff --git a/griptape/drivers/assistant/griptape_cloud_assistant_driver.py b/griptape/drivers/assistant/griptape_cloud_assistant_driver.py
index ada2aafe..1179f1ba 100644
--- a/griptape/drivers/assistant/griptape_cloud_assistant_driver.py
+++ b/griptape/drivers/assistant/griptape_cloud_assistant_driver.py
@@ -1,10 +1,10 @@
 from __future__ import annotations
 
+import json
 import logging
 import os
-import time
 import uuid
-from typing import Optional
+from typing import TYPE_CHECKING, Optional
 from urllib.parse import urljoin
 
 import requests
@@ -15,6 +15,9 @@ from griptape.configs.defaults_config import Defaults
 from griptape.drivers.assistant import BaseAssistantDriver
 from griptape.events import BaseEvent, EventBus
 
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
 logger = logging.getLogger(Defaults.logging_config.logger_name)
 
 
@@ -104,42 +107,33 @@ class GriptapeCloudAssistantDriver(BaseAssistantDriver):
         return response.json()["assistant_run_id"]
 
     def _get_run_result(self, assistant_run_id: str) -> TextArtifact:
-        events, next_offset = self._get_run_events(assistant_run_id)
-        attempts = 0
+        events = self._get_run_events(assistant_run_id)
         output = None
 
-        while output is None and attempts < self.max_attempts:
-            for event in events:
-                if event["origin"] == "ASSISTANT":
-                    event_payload = event["payload"]
-                    try:
-                        EventBus.publish_event(BaseEvent.from_dict(event_payload))
-                    except ValueError as e:
-                        logger.warning("Failed to deserialize event: %s", e)
-                    if event["type"] == "FinishStructureRunEvent":
-                        output = TextArtifact.from_dict(event_payload["output_task_output"])
-
-            if output is None and not events:
-                time.sleep(self.poll_interval)
-                attempts += 1
-            events, next_offset = self._get_run_events(assistant_run_id, offset=next_offset)
+        for event in events:
+            if event["origin"] == "ASSISTANT":
+                event_payload = event["payload"]
+                try:
+                    EventBus.publish_event(BaseEvent.from_dict(event_payload))
+                except ValueError as e:
+                    logger.warning("Failed to deserialize event: %s", e)
+                if event["type"] == "FinishStructureRunEvent":
+                    output = TextArtifact.from_dict(event_payload["output_task_output"])
 
         if output is None:
-            raise TimeoutError("The assistant run did not finish in time.")
+            raise ValueError("Output not found.")
 
         return output
 
-    def _get_run_events(self, assistant_run_id: str, offset: int = 0) -> tuple[list[dict], int]:
-        url = urljoin(self.base_url.strip("/"), f"/api/assistant-runs/{assistant_run_id}/events")
-        response = requests.get(url, headers=self.headers, params={"offset": offset})
-        response.raise_for_status()
-
-        response_json = response.json()
-
-        events = response_json.get("events", [])
-        next_offset = response_json.get("next_offset", 0)
-
-        return events, next_offset
+    def _get_run_events(self, assistant_run_id: str) -> Iterator[dict]:
+        url = urljoin(self.base_url.strip("/"), f"/api/assistant-runs/{assistant_run_id}/events/stream")
+        with requests.get(url, headers=self.headers, stream=True) as response:
+            response.raise_for_status()
+            for line in response.iter_lines():
+                if line:
+                    decoded_line = line.decode("utf-8")
+                    if decoded_line.startswith("data:"):
+                        yield json.loads(decoded_line.removeprefix("data:").strip())
 
     def _find_thread_by_alias(self, thread_alias: str) -> Optional[dict]:
         url = urljoin(self.base_url.strip("/"), "/api/threads")
diff --git a/griptape/drivers/structure_run/griptape_cloud_structure_run_driver.py b/griptape/drivers/structure_run/griptape_cloud_structure_run_driver.py
index 9855825a..db6b1a91 100644
--- a/griptape/drivers/structure_run/griptape_cloud_structure_run_driver.py
+++ b/griptape/drivers/structure_run/griptape_cloud_structure_run_driver.py
@@ -1,7 +1,8 @@
 from __future__ import annotations
 
+import json
 import logging
-import time
+from typing import TYPE_CHECKING
 from urllib.parse import urljoin
 
 import requests
@@ -12,6 +13,9 @@ from griptape.configs.defaults_config import Defaults
 from griptape.drivers.structure_run import BaseStructureRunDriver
 from griptape.events import BaseEvent, EventBus
 
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
 logger = logging.getLogger(Defaults.logging_config.logger_name)
 
 
@@ -52,43 +56,40 @@ class GriptapeCloudStructureRunDriver(BaseStructureRunDriver):
         return response_json["structure_run_id"]
 
     def _get_run_result(self, structure_run_id: str) -> BaseArtifact | InfoArtifact:
-        events, next_offset = self._get_run_events(structure_run_id)
-        attempts = 0
+        events = self._get_run_events(structure_run_id)
         output = None
 
-        while output is None and attempts < self.structure_run_max_wait_time_attempts:
-            for event in events:
-                event_type = event["type"]
-                event_payload = event.get("payload", {})
-                if event["origin"] == "USER":
-                    try:
-                        EventBus.publish_event(BaseEvent.from_dict(event_payload))
-                    except ValueError as e:
-                        logger.warning("Failed to deserialize event: %s", e)
-                    if event["type"] == "FinishStructureRunEvent":
-                        output = BaseArtifact.from_dict(event_payload["output_task_output"])
-                elif event["origin"] == "SYSTEM":
-                    if event_type == "StructureRunError":
-                        output = ErrorArtifact(event_payload["status_detail"]["error"])
-
-            if output is None and not events:
-                time.sleep(self.structure_run_wait_time_interval)
-                attempts += 1
-            events, next_offset = self._get_run_events(structure_run_id, offset=next_offset)
+        for event in events:
+            event_type = event["type"]
+            event_payload = event.get("payload", {})
+            if event["origin"] == "USER":
+                try:
+                    if "span_id" in event_payload:
+                        span_id = event_payload.pop("span_id")
+                        if "meta" in event_payload:
+                            event_payload["meta"]["span_id"] = span_id
+                        else:
+                            event_payload["meta"] = {"span_id": span_id}
+                    EventBus.publish_event(BaseEvent.from_dict(event_payload))
+                except ValueError as e:
+                    logger.warning("Failed to deserialize event: %s", e)
+                if event["type"] == "FinishStructureRunEvent":
+                    output = BaseArtifact.from_dict(event_payload["output_task_output"])
+            elif event["origin"] == "SYSTEM":
+                if event_type == "StructureRunError":
+                    output = ErrorArtifact(event_payload["status_detail"]["error"])
 
         if output is None:
-            raise TimeoutError("The structure run did not finish in time.")
+            raise ValueError("Output not found.")
 
         return output
 
-    def _get_run_events(self, structure_run_id: str, offset: int = 0) -> tuple[list[dict], int]:
-        url = urljoin(self.base_url.strip("/"), f"/api/structure-runs/{structure_run_id}/events")
-        response = requests.get(url, headers=self.headers, params={"offset": offset})
-        response.raise_for_status()
-
-        response_json = response.json()
-
-        events = response_json.get("events", [])
-        next_offset = response_json.get("next_offset", 0)
-
-        return events, next_offset
+    def _get_run_events(self, structure_run_id: str) -> Iterator[dict]:
+        url = urljoin(self.base_url.strip("/"), f"/api/structure-runs/{structure_run_id}/events/stream")
+        with requests.get(url, headers=self.headers, stream=True) as response:
+            response.raise_for_status()
+            for line in response.iter_lines():
+                if line:
+                    decoded_line = line.decode("utf-8")
+                    if decoded_line.startswith("data:"):
+                        yield json.loads(decoded_line.removeprefix("data:").strip())
diff --git a/griptape/mixins/activity_mixin.py b/griptape/mixins/activity_mixin.py
index 3db45104..cda8c88b 100644
--- a/griptape/mixins/activity_mixin.py
+++ b/griptape/mixins/activity_mixin.py
@@ -31,9 +31,6 @@ class ActivityMixin:
         if allowlist is None:
             return
 
-        if self.denylist is not None:
-            raise ValueError("can't have both allowlist and denylist specified")
-
         for activity_name in allowlist:
             self._validate_tool_activity(activity_name)
 
@@ -42,9 +39,6 @@ class ActivityMixin:
         if denylist is None:
             return
 
-        if self.allowlist is not None:
-            raise ValueError("can't have both allowlist and denylist specified")
-
         for activity_name in denylist:
             self._validate_tool_activity(activity_name)
 


Test patch:
diff --git a/tests/unit/drivers/assistant/test_griptape_cloud_assistant_driver.py b/tests/unit/drivers/assistant/test_griptape_cloud_assistant_driver.py
index dfe7ea4b..ad29c2fe 100644
--- a/tests/unit/drivers/assistant/test_griptape_cloud_assistant_driver.py
+++ b/tests/unit/drivers/assistant/test_griptape_cloud_assistant_driver.py
@@ -1,4 +1,5 @@
-from unittest.mock import ANY, call
+import json
+from unittest.mock import ANY, MagicMock, call
 
 import pytest
 import requests
@@ -43,51 +44,58 @@ class TestGriptapeCloudAssistantDriver:
         def request(*args, **kwargs):
             if "events" in args[0]:
                 mock_response = mocker.Mock()
-                mock_response.json.return_value = {
-                    "events": [
-                        {
-                            "origin": "ASSISTANT",
-                            "type": "FooBarEvent",
-                            "payload": {
+                mock_response.iter_lines.return_value = [
+                    *[
+                        f"data: {json.dumps(event)}".encode()
+                        for event in [
+                            {
+                                "origin": "ASSISTANT",
                                 "type": "FooBarEvent",
-                            },
-                        },
-                        {
-                            "origin": "ASSISTANT",
-                            "type": "FinishStructureRunEvent",
-                            "payload": {
-                                "type": "FinishStructureRunEvent",
-                                "output_task_input": {
-                                    "type": "TextArtifact",
-                                    "value": "foo bar",
-                                },
-                                "output_task_output": {
-                                    "type": "TextArtifact",
-                                    "value": "foo bar",
+                                "payload": {
+                                    "type": "FooBarEvent",
                                 },
                             },
-                        },
-                        {
-                            "origin": "ASSISTANT",
-                            "type": "FOO",
-                            "payload": {
+                            {
+                                "origin": "ASSISTANT",
                                 "type": "FinishStructureRunEvent",
-                                "output_task_input": {
-                                    "type": "TextArtifact",
-                                    "value": "foo bar",
+                                "payload": {
+                                    "type": "FinishStructureRunEvent",
+                                    "output_task_input": {
+                                        "type": "TextArtifact",
+                                        "value": "foo bar",
+                                    },
+                                    "output_task_output": {
+                                        "type": "TextArtifact",
+                                        "value": "foo bar",
+                                    },
                                 },
-                                "output_task_output": {
-                                    "type": "TextArtifact",
-                                    "value": "foo bar",
+                            },
+                            {
+                                "origin": "ASSISTANT",
+                                "type": "FOO",
+                                "payload": {
+                                    "type": "FinishStructureRunEvent",
+                                    "output_task_input": {
+                                        "type": "TextArtifact",
+                                        "value": "foo bar",
+                                    },
+                                    "output_task_output": {
+                                        "type": "TextArtifact",
+                                        "value": "foo bar",
+                                    },
                                 },
                             },
-                        },
-                        {
-                            "origin": "FOO",
-                        },
+                            {
+                                "origin": "FOO",
+                            },
+                        ]
                     ],
-                    "next_offset": 0,
-                }
+                    "",
+                    "title:",
+                ]
+
+                mock_response.__enter__ = MagicMock(return_value=mock_response)
+                mock_response.__exit__ = MagicMock()
 
                 return mock_response
             elif "threads" in args[0]:
@@ -114,15 +122,6 @@ class TestGriptapeCloudAssistantDriver:
             side_effect=request,
         )
 
-    @pytest.fixture()
-    def mock_requests_get_empty(self, mocker):
-        mock_response = mocker.Mock()
-        mock_response.json.return_value = {
-            "events": [],
-            "next_offset": 0,
-        }
-        return mocker.patch("requests.get", return_value=mock_response)
-
     @pytest.fixture()
     def driver(self):
         return GriptapeCloudAssistantDriver(
@@ -206,28 +205,23 @@ class TestGriptapeCloudAssistantDriver:
                 "additional_tool_ids": [],
             }
 
-    def test_timeout_run(self, driver, mocker):
+    def test_no_output_run(self, driver, mocker):
         mock_response = mocker.Mock()
-        mock_response.json.return_value = {
-            "events": [],
-            "next_offset": 0,
-        }
+        mock_response.__enter__ = MagicMock(return_value=mock_response)
+        mock_response.__exit__ = MagicMock()
+        mock_response.iter_lines.return_value = []
         mock_requests_get_empty = mocker.patch("requests.get", return_value=mock_response)
 
         driver.max_attempts = 1
-        with pytest.raises(TimeoutError):
+        with pytest.raises(ValueError, match="Output not found."):
             driver.run(TextArtifact("foo bar"))
 
         expected_calls = [
             call(
-                "https://cloud-foo.griptape.ai/api/assistant-runs/1/events",
-                params={"offset": 0},
-                headers={"Authorization": "Bearer foo bar"},
-            ),
-            call(
-                "https://cloud-foo.griptape.ai/api/assistant-runs/1/events",
-                params={"offset": 0},
+                "https://cloud-foo.griptape.ai/api/assistant-runs/1/events/stream",
                 headers={"Authorization": "Bearer foo bar"},
+                stream=True,
             ),
         ]
+
         mock_requests_get_empty.assert_has_calls(expected_calls)
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
index 993c5af6..bbad6afe 100644
--- 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
@@ -1,4 +1,5 @@
-from unittest.mock import Mock
+import json
+from unittest.mock import MagicMock, Mock
 
 import pytest
 
@@ -13,43 +14,64 @@ class TestGriptapeCloudStructureRunDriver:
     @pytest.fixture(autouse=True)
     def mock_requests_get(self, mocker):
         mock_response = mocker.Mock()
-        mock_response.json.return_value = {
-            "events": [
-                {
-                    "origin": "USER",
-                    "type": "FooBarEvent",
-                    "payload": {
+        mock_response.iter_lines.return_value = [
+            *[
+                f"data: {json.dumps(event)}".encode()
+                for event in [
+                    {
+                        "origin": "USER",
                         "type": "FooBarEvent",
+                        "payload": {
+                            "type": "FooBarEvent",
+                            "span_id": "1",
+                        },
                     },
-                },
-                {
-                    "origin": "USER",
-                    "type": "FinishStructureRunEvent",
-                    "payload": {
+                    {
+                        "origin": "USER",
                         "type": "FinishStructureRunEvent",
-                        "output_task_input": {
-                            "type": "TextArtifact",
-                            "value": "foo bar",
+                        "payload": {
+                            "type": "FinishStructureRunEvent",
+                            "span_id": "1",
+                            "output_task_input": {
+                                "type": "TextArtifact",
+                                "value": "foo bar",
+                            },
+                            "output_task_output": {
+                                "type": "TextArtifact",
+                                "value": "foo bar",
+                            },
+                            "meta": {
+                                "foo": "bar",
+                            },
                         },
-                        "output_task_output": {
-                            "type": "TextArtifact",
-                            "value": "foo bar",
+                    },
+                    {"origin": "FOO", "type": "BAR"},
+                    {
+                        "origin": "SYSTEM",
+                        "type": "StructureRunError",
+                        "payload": {
+                            "status_detail": {
+                                "error": "foo bar",
+                            },
                         },
                     },
-                },
-                {"origin": "FOO", "type": "BAR"},
-                {
-                    "origin": "SYSTEM",
-                    "type": "StructureRunError",
-                    "payload": {
-                        "status_detail": {
-                            "error": "foo bar",
+                    {
+                        "origin": "SYSTEM",
+                        "type": "StructureRunCompleted",
+                        "payload": {
+                            "status": "SUCCEEDED",
+                            "started_at": "2024-09-11T23:15:47",
+                            "completed_at": "2024-09-11T23:15:50",
+                            "status_detail": {"reason": "Completed", "message": None, "exit_code": 0},
                         },
                     },
-                },
+                ]
             ],
-            "next_offset": 0,
-        }
+            "",
+            "title:",
+        ]
+        mock_response.__enter__ = MagicMock(return_value=mock_response)
+        mock_response.__exit__ = MagicMock()
 
         return mocker.patch(
             "requests.get",
@@ -84,6 +106,10 @@ class TestGriptapeCloudStructureRunDriver:
         events = mock_on_event.call_args[0]
         assert len(events) == 1
         assert events[0].type == "FinishStructureRunEvent"
+        assert events[0].meta == {
+            "foo": "bar",
+            "span_id": "1",
+        }
 
         assert isinstance(result, ErrorArtifact)
         assert result.value == "foo bar"
@@ -95,16 +121,11 @@ class TestGriptapeCloudStructureRunDriver:
         assert result.value == "Run started successfully"
 
     def test_run_timeout(self, driver, mocker):
-        mocker.patch(
-            "requests.get",
-            return_value=Mock(
-                json=Mock(
-                    return_value={
-                        "events": [],
-                    }
-                )
-            ),
-        )
+        mock_response = mocker.Mock()
+        mock_response.__enter__ = MagicMock(return_value=mock_response)
+        mock_response.__exit__ = MagicMock()
+        mock_response.iter_lines.return_value = []
+        mocker.patch("requests.get", return_value=mock_response)
 
-        with pytest.raises(TimeoutError):
+        with pytest.raises(ValueError, match="Output not found."):
             driver.run(TextArtifact("foo bar"))
diff --git a/tests/unit/mixins/test_activity_mixin.py b/tests/unit/mixins/test_activity_mixin.py
index bef7c118..d1b810b4 100644
--- a/tests/unit/mixins/test_activity_mixin.py
+++ b/tests/unit/mixins/test_activity_mixin.py
@@ -43,17 +43,33 @@ class TestActivityMixin:
 
     def test_allowlist_and_denylist_validation(self):
         with pytest.raises(ValueError):
-            MockTool(test_field="hello", test_int=5, allowlist=[], denylist=[])
+            MockTool(test_field="hello", test_int=5, allowlist=["not_an_activity"], denylist=[])
 
-    def test_allowlist(self):
-        tool = MockTool(test_field="hello", test_int=5, allowlist=["test"])
-
-        assert len(tool.activities()) == 1
-
-    def test_denylist(self):
-        tool = MockTool(test_field="hello", test_int=5, denylist=["test"])
+    @pytest.mark.parametrize(
+        ("allowlist", "denylist", "expected_activities"),
+        [
+            (["test"], None, ["test"]),
+            (["test"], ["test"], []),
+            (
+                None,
+                ["test"],
+                [
+                    "test_callable_schema",
+                    "test_error",
+                    "test_exception",
+                    "test_list_output",
+                    "test_no_schema",
+                    "test_str_output",
+                    "test_without_default_memory",
+                ],
+            ),
+            ([], ["test"], []),
+        ],
+    )
+    def test_allowdenylist(self, allowlist, denylist, expected_activities):
+        tool = MockTool(test_field="hello", test_int=5, allowlist=allowlist, denylist=denylist)
 
-        assert len(tool.activities()) == 7
+        assert [getattr(activity, "name") for activity in tool.activities()] == expected_activities
 
     def test_invalid_allowlist(self):
         with pytest.raises(ValueError):