Automatically create Thread in Griptape Cloud Assistant Driver (#1582)

Author: collindutter

Diff for: CHANGELOG.md

@@ -17,6 +17,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 - Task bitshift operators can now take a list of Tasks.
+- `GriptapeCloudAssistantDriver` and `OpenAiAssistantDriver` now automatically create a new Thread if one is not provided. Can be disabled with `auto_create_thread=False`.
+- `GriptapeCloudAssistantDriver` and `OpenAiAssistantDriver` now return metadata (`thread_id`) on the response Artifact.
+- `GriptapeCloudAssistantDriver` now accepts a `thread_alias` parameter for fetching a Thread by alias, creating one if it doesn't exist.
+- `EvalEngine` to use structured output when generating evaluation steps.
+- `GriptapeCloudVectorStoreDriver.query()` updated to non-deprecated Griptape Cloud API shape.
 
 ### Fixed
 
@@ -34,10 +39,6 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Deprecated
 
 - `FuturesExecutorMixin.futures_executor`. Use `FuturesExecutorMixin.create_futures_executor` instead.
-### Changed
-
-- `EvalEngine` to use structured output when generating evaluation steps.
-- `GriptapeCloudVectorStoreDriver.query()` updated to non-deprecated Griptape Cloud API shape.
 
 ## [1.1.1] - 2025-01-03
 

Diff for: griptape/drivers/assistant/griptape_cloud_assistant_driver.py

@@ -3,6 +3,7 @@
 import logging
 import os
 import time
+import uuid
 from typing import Optional
 from urllib.parse import urljoin
 
@@ -30,6 +31,7 @@ class GriptapeCloudAssistantDriver(BaseAssistantDriver):
     input: Optional[str] = field(default=None, kw_only=True)
     assistant_id: str = field(kw_only=True)
     thread_id: Optional[str] = field(default=None, kw_only=True)
+    thread_alias: Optional[str] = field(default=None, kw_only=True)
     ruleset_ids: Optional[list[str]] = field(default=None, kw_only=True)
     additional_ruleset_ids: list[str] = field(factory=list, kw_only=True)
     knowledge_base_ids: Optional[list[str]] = field(default=None, kw_only=True)
@@ -41,8 +43,41 @@ class GriptapeCloudAssistantDriver(BaseAssistantDriver):
     stream: bool = field(default=False, kw_only=True)
     poll_interval: int = field(default=1, kw_only=True)
     max_attempts: int = field(default=20, kw_only=True)
+    auto_create_thread: bool = field(default=True, kw_only=True)
 
     def try_run(self, *args: BaseArtifact) -> BaseArtifact | InfoArtifact:
+        if self.thread_id is None and self.auto_create_thread:
+            self._create_or_find_thread(self.thread_alias)
+        assistant_run_id = self._create_run(*args)
+        run_result = self._get_run_result(assistant_run_id)
+
+        run_result.meta.update(
+            {"assistant_id": self.assistant_id, "assistant_run_id": assistant_run_id, "thread_id": self.thread_id}
+        )
+
+        return run_result
+
+    def _create_or_find_thread(self, thread_alias: Optional[str] = None) -> None:
+        if thread_alias is None:
+            self.thread_id = self._create_thread()
+        else:
+            thread = self._find_thread_by_alias(thread_alias)
+
+            if thread is None:
+                self.thread_id = self._create_thread(thread_alias)
+
+    def _create_thread(self, thread_alias: Optional[str] = None) -> str:
+        url = urljoin(self.base_url.strip("/"), "/api/threads")
+
+        body = {"name": uuid.uuid4().hex}
+        if thread_alias is not None:
+            body["alias"] = thread_alias
+
+        response = requests.post(url, json=body, headers=self.headers)
+        response.raise_for_status()
+        return response.json()["thread_id"]
+
+    def _create_run(self, *args: BaseArtifact) -> str:
         url = urljoin(self.base_url.strip("/"), f"/api/assistants/{self.assistant_id}/runs")
 
         response = requests.post(
@@ -64,9 +99,7 @@ def try_run(self, *args: BaseArtifact) -> BaseArtifact | InfoArtifact:
             headers=self.headers,
         )
         response.raise_for_status()
-        response_json = response.json()
-
-        return self._get_run_result(response_json["assistant_run_id"])
+        return response.json()["assistant_run_id"]
 
     def _get_run_result(self, assistant_run_id: str) -> BaseArtifact | InfoArtifact:
         events, next_offset = self._get_run_events(assistant_run_id)
@@ -105,3 +138,12 @@ def _get_run_events(self, assistant_run_id: str, offset: int = 0) -> tuple[list[
         next_offset = response_json.get("next_offset", 0)
 
         return events, next_offset
+
+    def _find_thread_by_alias(self, thread_alias: str) -> Optional[dict]:
+        url = urljoin(self.base_url.strip("/"), "/api/threads")
+        response = requests.get(url, params={"alias": thread_alias}, headers=self.headers)
+        response.raise_for_status()
+
+        threads = response.json()["threads"]
+
+        return next((thread for thread in threads if thread["alias"] == thread_alias), None)

Diff for: griptape/drivers/assistant/openai_assistant_driver.py

@@ -39,11 +39,12 @@ def on_tool_call_delta(self, delta: ToolCallDelta, snapshot: ToolCall) -> None:
     base_url: Optional[str] = field(default=None, kw_only=True, metadata={"serializable": True})
     api_key: Optional[str] = field(default=None, kw_only=True, metadata={"serializable": False})
     organization: Optional[str] = field(default=None, kw_only=True, metadata={"serializable": True})
-    thread_id: Optional[str] = field(kw_only=True)
+    thread_id: Optional[str] = field(default=None, kw_only=True)
     assistant_id: str = field(kw_only=True)
     event_handler: AssistantEventHandler = field(
         default=Factory(lambda: OpenAiAssistantDriver.EventHandler()), kw_only=True, metadata={"serializable": False}
     )
+    auto_create_thread: bool = field(default=True, kw_only=True)
 
     _client: openai.OpenAI = field(default=None, kw_only=True, alias="client", metadata={"serializable": False})
 
@@ -56,8 +57,17 @@ def client(self) -> openai.OpenAI:
         )
 
     def try_run(self, *args: BaseArtifact) -> BaseArtifact | InfoArtifact:
+        if self.thread_id is None and self.auto_create_thread:
+            self.thread_id = self.client.beta.threads.create().id
+        response = self._create_run(*args)
+
+        response.meta.update({"assistant_id": self.assistant_id, "thread_id": self.thread_id})
+
+        return response
+
+    def _create_run(self, *args: BaseArtifact) -> TextArtifact:
         content = "\n".join(arg.value for arg in args)
-        self.client.beta.threads.messages.create(thread_id=self.thread_id, role="user", content=content)
+        message_id = self.client.beta.threads.messages.create(thread_id=self.thread_id, role="user", content=content)
         with self.client.beta.threads.runs.stream(
             thread_id=self.thread_id,
             assistant_id=self.assistant_id,
@@ -71,4 +81,9 @@ def try_run(self, *args: BaseArtifact) -> BaseArtifact | InfoArtifact:
                 message_contents.append("".join(content.text.value for content in message.content))
             message_text = "\n".join(message_contents)
 
-            return TextArtifact(message_text)
+            response = TextArtifact(message_text)
+
+            response.meta.update(
+                {"assistant_id": self.assistant_id, "thread_id": self.thread_id, "message_id": message_id}
+            )
+            return response