feat(utils): add griptape cloud context util

WIP

WIP

Author: collindutter

CHANGELOG.md

@@ -31,7 +31,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Support for new serialization metadata, `serialization_key` and `deserialization_key` for more granular control over serialization. 
 - Support for OpenAi reasoning models, `o1` and `o3`.
 - Support for enums in `GriptapeCloudToolTool`.
-- `griptape.utils.griptape_cloud_utils.configure_events` for automatically configuring the `EventBus` for use with Griptape Cloud Managed Structures.
+- `griptape.utils.griptape_cloud_utils.GriptapeCloud` for automatically configuring the `EventBus` for use with Griptape Cloud Managed Structures.
 
 ### Changed
 

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 GriptapeCloud
 
 
 def minify_json(value: str) -> str:
@@ -58,4 +59,5 @@ def minify_json(value: str) -> str:
     "with_contextvars",
     "build_strict_schema",
     "resolve_refs",
+    "GriptapeCloud",
 ]

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 GriptapeCloud:
+    """Utility for working with Griptape Cloud.
+
+    Args:
+        _event_listener (EventListener): Event Listener to use. Defaults to an EventListener with a GriptapeCloudEventListenerDriver.
+        _observability (Observability): Observability to use. Defaults to an Observability with a GriptapeCloudObservabilityDriver.
+        observe (bool): 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) -> GriptapeCloud:
+        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)

griptape/utils/griptape_cloud_utils.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/drivers/event_listener/test_griptape_cloud_event_listener_driver.py

@@ -1,16 +1,58 @@
 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.utils.griptape_cloud_utils import configure_events
+from griptape.observability.observability import Observability
+from griptape.utils import GriptapeCloud
+from tests.mocks.mock_event_listener_driver import MockEventListenerDriver
 
 
 class TestGriptapeCloudUtils:
-    def test_configure_events(self):
+    @pytest.fixture(autouse=True)
+    def set_up_environment(self):
         os.environ["GT_CLOUD_API_KEY"] = "foo"
         os.environ["GT_CLOUD_STRUCTURE_RUN_ID"] = "bar"
 
-        configure_events()
+        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 GriptapeCloud(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 = GriptapeCloud()
+
+        assert context.in_managed_environment
+
+    def test_structure_run_id(self):
+        context = GriptapeCloud()
+
+        assert context.structure_run_id == "bar"
+
+    def test_output(self):
+        with GriptapeCloud(event_listener=EventListener(event_listener_driver=MockEventListenerDriver())) as context:
+            context.output = "foo"
+            assert context.output.value == "foo"
 
-        assert len(EventBus.event_listeners) == 1
-        assert isinstance(EventBus.event_listeners[0].event_listener_driver, GriptapeCloudEventListenerDriver)
+            output = TextArtifact("bar")
+            context.output = output
+            assert context.output is output