From 10c0170af606654ccee116ecd44eb159ca510bdb Mon Sep 17 00:00:00 2001 From: Matt Vallillo Date: Thu, 29 Aug 2024 14:34:14 -0400 Subject: [PATCH] Update GriptapeCloudEventListenerDriver (#1115) --- .../griptape_cloud_event_listener_driver.py | 24 ++++++++++++------- tests/mocks/mock_event.py | 2 +- ...st_griptape_cloud_event_listener_driver.py | 8 +++---- 3 files changed, 21 insertions(+), 13 deletions(-) diff --git a/griptape/drivers/event_listener/griptape_cloud_event_listener_driver.py b/griptape/drivers/event_listener/griptape_cloud_event_listener_driver.py index 98f52b914..3e06eaa88 100644 --- a/griptape/drivers/event_listener/griptape_cloud_event_listener_driver.py +++ b/griptape/drivers/event_listener/griptape_cloud_event_listener_driver.py @@ -54,13 +54,21 @@ def publish_event(self, event: BaseEvent | dict, *, flush: bool = False) -> None super().publish_event(event_payload, flush=flush) def try_publish_event_payload(self, event_payload: dict) -> None: - url = urljoin(self.base_url.strip("/"), f"/api/structure-runs/{self.structure_run_id}/events") - - response = requests.post(url=url, json=event_payload, headers=self.headers) - response.raise_for_status() + self._post_event(self._get_event_request(event_payload)) def try_publish_event_payload_batch(self, event_payload_batch: list[dict]) -> None: - url = urljoin(self.base_url.strip("/"), f"/api/structure-runs/{self.structure_run_id}/events") - - response = requests.post(url=url, json=event_payload_batch, headers=self.headers) - response.raise_for_status() + self._post_event([self._get_event_request(event_payload) for event_payload in event_payload_batch]) + + def _get_event_request(self, event_payload: dict) -> dict: + return { + "payload": event_payload, + "timestamp": event_payload["timestamp"], + "type": event_payload["type"], + } + + def _post_event(self, json: list[dict] | dict) -> None: + requests.post( + url=urljoin(self.base_url.strip("/"), f"/api/structure-runs/{self.structure_run_id}/events"), + json=json, + headers=self.headers, + ).raise_for_status() diff --git a/tests/mocks/mock_event.py b/tests/mocks/mock_event.py index 2b9d9ade3..a8737f47c 100644 --- a/tests/mocks/mock_event.py +++ b/tests/mocks/mock_event.py @@ -3,4 +3,4 @@ class MockEvent(BaseEvent): def to_dict(self) -> dict: - return {"timestamp": self.timestamp, "id": self.id} + return {"timestamp": self.timestamp, "id": self.id, "meta": self.meta, "type": self.__class__.__name__} diff --git a/tests/unit/drivers/event_listener/test_griptape_cloud_event_listener_driver.py b/tests/unit/drivers/event_listener/test_griptape_cloud_event_listener_driver.py index 0bf298870..441589774 100644 --- a/tests/unit/drivers/event_listener/test_griptape_cloud_event_listener_driver.py +++ b/tests/unit/drivers/event_listener/test_griptape_cloud_event_listener_driver.py @@ -49,7 +49,7 @@ def test_publish_event_without_span_id(self, mock_post, driver): mock_post.assert_called_with( url="https://cloud123.griptape.ai/api/structure-runs/bar baz/events", - json=[event.to_dict()], + json=[driver._get_event_request(event.to_dict())], headers={"Authorization": "Bearer foo bar"}, ) @@ -63,7 +63,7 @@ def test_publish_event_with_span_id(self, mock_post, driver): mock_post.assert_called_with( url="https://cloud123.griptape.ai/api/structure-runs/bar baz/events", - json=[{**event.to_dict(), "span_id": "test"}], + json=[driver._get_event_request({**event.to_dict(), "span_id": "test"})], headers={"Authorization": "Bearer foo bar"}, ) @@ -73,7 +73,7 @@ def test_try_publish_event_payload(self, mock_post, driver): mock_post.assert_called_once_with( url="https://cloud123.griptape.ai/api/structure-runs/bar baz/events", - json=event.to_dict(), + json=driver._get_event_request(event.to_dict()), headers={"Authorization": "Bearer foo bar"}, ) @@ -84,6 +84,6 @@ def try_publish_event_payload_batch(self, mock_post, driver): mock_post.assert_called_with( url="https://cloud123.griptape.ai/api/structure-runs/bar baz/events", - json=event.to_dict(), + json=driver._get_event_request(event.to_dict()), headers={"Authorization": "Bearer foo bar"}, )