feat(utils): add utility function for configuring Event Bus with GTC (#1651)

Author: collindutter

CHANGELOG.md

@@ -32,6 +32,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Support for OpenAi reasoning models, `o1` and `o3`.
 - Support for enums in `GriptapeCloudToolTool`.
 - `LocalRerankDriver` for reranking locally.
+- `griptape.utils.griptape_cloud.GriptapeCloudStructure` for automatically configuring Cloud-specific Drivers when in the Griptape Cloud Structures Runtime.
 
 ### Changed
 

griptape/drivers/event_listener/griptape_cloud_event_listener_driver.py

@@ -26,7 +26,7 @@ class GriptapeCloudEventListenerDriver(BaseEventListenerDriver):
         default=Factory(lambda: os.getenv("GT_CLOUD_BASE_URL", "https://cloud.griptape.ai")),
         kw_only=True,
     )
-    api_key: str = field(default=Factory(lambda: os.environ["GT_CLOUD_API_KEY"]), kw_only=True)
+    api_key: Optional[str] = field(default=Factory(lambda: os.getenv("GT_CLOUD_API_KEY")), kw_only=True)
     headers: dict = field(
         default=Factory(lambda self: {"Authorization": f"Bearer {self.api_key}"}, takes_self=True),
         kw_only=True,
@@ -36,12 +36,22 @@ class GriptapeCloudEventListenerDriver(BaseEventListenerDriver):
     )
 
     @structure_run_id.validator  # pyright: ignore[reportAttributeAccessIssue, reportOptionalMemberAccess]
-    def validate_run_id(self, _: Attribute, structure_run_id: str) -> None:
+    def validate_run_id(self, _: Attribute, structure_run_id: Optional[str]) -> None:
         if structure_run_id is None:
             raise ValueError(
                 "structure_run_id must be set either in the constructor or as an environment variable (GT_CLOUD_STRUCTURE_RUN_ID).",
             )
 
+    @api_key.validator  # pyright: ignore[reportAttributeAccessIssue, reportOptionalMemberAccess]
+    def validate_api_key(self, _: Attribute, api_key: Optional[str]) -> None:
+        if api_key is None:
+            raise ValueError(
+                "No value was found for the 'GT_CLOUD_API_KEY' environment variable. "
+                "This environment variable is required when running in Griptape Cloud for authorization. "
+                "You can generate a Griptape Cloud API Key by visiting https://cloud.griptape.ai/keys . "
+                "Specify it as an environment variable when creating a Managed Structure in Griptape Cloud."
+            )
+
     def publish_event(self, event: BaseEvent | dict) -> None:
         from griptape.observability.observability import Observability
 

griptape/events/finish_structure_run_event.py

@@ -13,5 +13,5 @@
 @define
 class FinishStructureRunEvent(BaseEvent):
     structure_id: Optional[str] = field(kw_only=True, default=None, metadata={"serializable": True})
-    output_task_input: BaseArtifact = field(kw_only=True, metadata={"serializable": True})
-    output_task_output: Optional[BaseArtifact] = field(kw_only=True, metadata={"serializable": True})
+    output_task_input: Optional[BaseArtifact] = field(kw_only=True, default=None, metadata={"serializable": True})
+    output_task_output: Optional[BaseArtifact] = field(kw_only=True, default=None, metadata={"serializable": True})

griptape/utils/__init__.py

@@ -24,6 +24,7 @@
 from .file_utils import get_mime_type
 from .contextvars_utils import with_contextvars
 from .json_schema_utils import build_strict_schema, resolve_refs
+from .griptape_cloud import GriptapeCloudStructure
 
 
 def minify_json(value: str) -> str:
@@ -58,4 +59,5 @@ def minify_json(value: str) -> str:
     "with_contextvars",
     "build_strict_schema",
     "resolve_refs",
+    "GriptapeCloudStructure",
 ]

griptape/utils/griptape_cloud.py

@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+import os
+from typing import TYPE_CHECKING, Any, Optional, TypeVar
+
+from attrs import define, field
+from typing_extensions import ParamSpec
+
+from griptape.artifacts import BaseArtifact, GenericArtifact
+from griptape.events import EventBus, EventListener, FinishStructureRunEvent
+from griptape.utils.decorators import lazy_property
+
+P = ParamSpec("P")
+T = TypeVar("T")
+
+if TYPE_CHECKING:
+    from types import TracebackType
+
+    from griptape.observability.observability import Observability
+
+
+@define()
+class GriptapeCloudStructure:
+    """Utility for working with Griptape Cloud Structures.
+
+    Attributes:
+        _event_listener: Event Listener to use. Defaults to an EventListener with a GriptapeCloudEventListenerDriver.
+        _observability: Observability to use. Defaults to an Observability with a GriptapeCloudObservabilityDriver.
+        observe: Whether to enable observability. Enabling requires the `drivers-observability-griptape-cloud` extra.
+    """
+
+    _event_listener: EventListener = field(default=None, kw_only=True, alias="event_listener")
+    _observability: Observability = field(default=None, kw_only=True, alias="observability")
+    observe: bool = field(default=False, kw_only=True)
+    _output: BaseArtifact = field(default=None, init=False)
+
+    @lazy_property()
+    def event_listener(self) -> EventListener:
+        from griptape.drivers.event_listener.griptape_cloud import GriptapeCloudEventListenerDriver
+
+        return EventListener(event_listener_driver=GriptapeCloudEventListenerDriver())
+
+    @lazy_property()
+    def observability(self) -> Observability:
+        from griptape.drivers.observability.griptape_cloud import GriptapeCloudObservabilityDriver
+        from griptape.observability.observability import Observability
+
+        return Observability(observability_driver=GriptapeCloudObservabilityDriver())
+
+    @property
+    def output(self) -> BaseArtifact:
+        return self._output
+
+    @output.setter
+    def output(self, value: BaseArtifact | Any) -> None:
+        if not isinstance(value, BaseArtifact):
+            self._output = GenericArtifact(value)
+        else:
+            self._output = value
+
+    @property
+    def structure_run_id(self) -> str:
+        return os.environ["GT_CLOUD_STRUCTURE_RUN_ID"]
+
+    @property
+    def in_managed_environment(self) -> bool:
+        return "GT_CLOUD_STRUCTURE_RUN_ID" in os.environ
+
+    def __enter__(self) -> GriptapeCloudStructure:
+        from griptape.observability.observability import Observability
+
+        if self.in_managed_environment:
+            EventBus.add_event_listener(self.event_listener)
+
+            if self.observe:
+                Observability.set_global_driver(self.observability.observability_driver)
+                self.observability.observability_driver.__enter__()
+
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[type[BaseException]],
+        exc_value: Optional[BaseException],
+        exc_traceback: Optional[TracebackType],
+    ) -> None:
+        from griptape.observability.observability import Observability
+
+        if self.in_managed_environment:
+            if self.output is not None:
+                EventBus.publish_event(FinishStructureRunEvent(output_task_output=self.output), flush=True)
+            EventBus.remove_event_listener(self.event_listener)
+
+            if self.observe:
+                Observability.set_global_driver(None)
+                self.observability.observability_driver.__exit__(exc_type, exc_value, exc_traceback)

tests/unit/drivers/event_listener/test_griptape_cloud_event_listener_driver.py

@@ -83,7 +83,17 @@ def test_try_publish_event_payload(self, mock_post, driver):
             headers={"Authorization": "Bearer foo bar"},
         )
 
-    def try_publish_event_payload_batch(self, mock_post, driver):
+    def test_validate_api_key(self):
+        with pytest.raises(ValueError, match="No value was found"):
+            GriptapeCloudEventListenerDriver()
+
+    def test_validate_run_id(self):
+        os.environ["GT_CLOUD_API_KEY"] = "foo bar"
+
+        with pytest.raises(ValueError, match="structure_run_id must be set"):
+            GriptapeCloudEventListenerDriver()
+
+    def test_try_publish_event_payload_batch(self, mock_post, driver):
         for _ in range(3):
             event = MockEvent()
             driver.try_publish_event_payload(event.to_dict())

tests/unit/utils/test_griptape_cloud_utils.py

@@ -0,0 +1,60 @@
+import os
+
+import pytest
+
+from griptape.artifacts.text_artifact import TextArtifact
+from griptape.drivers.event_listener.griptape_cloud import GriptapeCloudEventListenerDriver
+from griptape.events import EventListener
+from griptape.events.event_bus import EventBus
+from griptape.observability.observability import Observability
+from griptape.utils import GriptapeCloudStructure
+from tests.mocks.mock_event_listener_driver import MockEventListenerDriver
+
+
+class TestGriptapeCloudUtils:
+    @pytest.fixture(autouse=True)
+    def set_up_environment(self):
+        os.environ["GT_CLOUD_API_KEY"] = "foo"
+        os.environ["GT_CLOUD_STRUCTURE_RUN_ID"] = "bar"
+
+        yield
+
+        del os.environ["GT_CLOUD_API_KEY"]
+        del os.environ["GT_CLOUD_STRUCTURE_RUN_ID"]
+
+    @pytest.mark.parametrize("observe", [True, False])
+    def test_context_manager(self, observe):
+        with GriptapeCloudStructure(observe=observe) as context:
+            assert context.in_managed_environment
+
+            assert len(EventBus.event_listeners) == 1
+            assert isinstance(EventBus.event_listeners[0].event_listener_driver, GriptapeCloudEventListenerDriver)
+
+            if observe:
+                assert Observability.get_global_driver() == context.observability.observability_driver
+            else:
+                assert Observability.get_global_driver() is None
+
+        assert len(EventBus.event_listeners) == 0
+        assert Observability.get_global_driver() is None
+
+    def test_in_managed_environment(self):
+        context = GriptapeCloudStructure()
+
+        assert context.in_managed_environment
+
+    def test_structure_run_id(self):
+        context = GriptapeCloudStructure()
+
+        assert context.structure_run_id == "bar"
+
+    def test_output(self):
+        with GriptapeCloudStructure(
+            event_listener=EventListener(event_listener_driver=MockEventListenerDriver())
+        ) as context:
+            context.output = "foo"
+            assert context.output.value == "foo"
+
+            output = TextArtifact("bar")
+            context.output = output
+            assert context.output is output