Improve Assistants API

src/marvin/beta/assistants/assistants.py

@@ -1,6 +1,6 @@
 from typing import TYPE_CHECKING, Callable, Optional, Union
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, PrivateAttr
 
 import marvin.utilities.tools
 from marvin.tools.assistants import AssistantTool
@@ -11,7 +11,7 @@
     run_sync,
 )
 from marvin.utilities.logging import get_logger
-from marvin.utilities.openai import get_client
+from marvin.utilities.openai import get_openai_client
 
 from .threads import Thread
 
@@ -29,6 +29,8 @@ class Assistant(BaseModel, ExposeSyncMethodsMixin):
     tools: list[Union[AssistantTool, Callable]] = []
     file_ids: list[str] = []
     metadata: dict[str, str] = {}
+    # context level tracks nested assistant contexts
+    _context_level: int = PrivateAttr(0)
 
     default_thread: Thread = Field(
         default_factory=Thread,
@@ -57,47 +59,61 @@ async def say_async(
         self,
         message: str,
         file_paths: Optional[list[str]] = None,
+        thread: Optional[Thread] = None,
         **run_kwargs,
     ) -> "Run":
         """
         A convenience method for adding a user message to the assistant's
         default thread, running the assistant, and returning the assistant's
         messages.
         """
+        thread = thread or self.default_thread
+
+        last_message = await thread.get_messages_async(limit=1)
+        if last_message:
+            last_msg_id = last_message[0].id
+        else:
+            last_msg_id = None
+
+        # post the message
         if message:
-            await self.default_thread.add_async(message, file_paths=file_paths)
+            await thread.add_async(message, file_paths=file_paths)
 
-        run = await self.default_thread.run_async(
-            assistant=self,
-            **run_kwargs,
-        )
-        return run
+        # run the thread
+        async with self:
+            await thread.run_async(assistant=self, **run_kwargs)
+
+        # load all messages, including the user message
+        response_messages = await thread.get_messages_async(after_message=last_msg_id)
+        return response_messages
 
     def __enter__(self):
-        self.create()
-        return self
+        return run_sync(self.__aenter__())
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        self.delete()
-        # If an exception has occurred, you might want to handle it or pass it through
-        # Returning False here will re-raise any exception that occurred in the context
-        return False
+        return run_sync(self.__aexit__(exc_type, exc_val, exc_tb))
 
     async def __aenter__(self):
-        await self.create_async()
+        self._context_level += 1
+        # if this is the outermost context and no ID is set, create the assistant
+        if self.id is None and self._context_level == 1:
+            await self.create_async()
         return self
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):
-        await self.delete_async()
-        # If an exception has occurred, you might want to handle it or pass it through
-        # Returning False here will re-raise any exception that occurred in the context
+        # If this is the outermost context, delete the assistant
+        if self._context_level == 1:
+            await self.delete_async()
+        self._context_level -= 1
         return False
 
     @expose_sync_method("create")
     async def create_async(self):
         if self.id is not None:
-            raise ValueError("Assistant has already been created.")
-        client = get_client()
+            raise ValueError(
+                "Assistant has an ID and has already been created in the OpenAI API."
+            )
+        client = get_openai_client()
         response = await client.beta.assistants.create(
             **self.model_dump(
                 include={"name", "model", "metadata", "file_ids", "metadata"}
@@ -106,25 +122,24 @@ async def create_async(self):
             instructions=self.get_instructions(),
         )
         self.id = response.id
-        self.clear_default_thread()
 
     @expose_sync_method("delete")
     async def delete_async(self):
         if not self.id:
-            raise ValueError("Assistant has not been created.")
-        client = get_client()
+            raise ValueError("Assistant has no ID and doesn't exist in the OpenAI API.")
+        client = get_openai_client()
         await client.beta.assistants.delete(assistant_id=self.id)
         self.id = None
 
     @classmethod
-    def load(cls, assistant_id: str):
-        return run_sync(cls.load_async(assistant_id))
+    def load(cls, assistant_id: str, **kwargs):
+        return run_sync(cls.load_async(assistant_id, **kwargs))
 
     @classmethod
-    async def load_async(cls, assistant_id: str):
-        client = get_client()
+    async def load_async(cls, assistant_id: str, **kwargs):
+        client = get_openai_client()
         response = await client.beta.assistants.retrieve(assistant_id=assistant_id)
-        return cls.model_validate(response)
+        return cls(**(response.model_dump() | kwargs))
 
     def chat(self, thread: Thread = None):
         if thread is None:

src/marvin/beta/assistants/runs.py

@@ -10,7 +10,7 @@
 from marvin.types import Tool
 from marvin.utilities.asyncio import ExposeSyncMethodsMixin, expose_sync_method
 from marvin.utilities.logging import get_logger
-from marvin.utilities.openai import get_client
+from marvin.utilities.openai import get_openai_client
 
 from .assistants import Assistant
 from .threads import Thread
@@ -54,20 +54,20 @@ def format_tools(cls, tools: Union[None, list[Union[Tool, Callable]]]):
 
     @expose_sync_method("refresh")
     async def refresh_async(self):
-        client = get_client()
+        client = get_openai_client()
         self.run = await client.beta.threads.runs.retrieve(
             run_id=self.run.id, thread_id=self.thread.id
         )
 
     @expose_sync_method("cancel")
     async def cancel_async(self):
-        client = get_client()
+        client = get_openai_client()
         await client.beta.threads.runs.cancel(
             run_id=self.run.id, thread_id=self.thread.id
         )
 
     async def _handle_step_requires_action(self):
-        client = get_client()
+        client = get_openai_client()
         if self.run.status != "requires_action":
             return
         if self.run.required_action.type == "submit_tool_outputs":
@@ -117,7 +117,7 @@ def get_tools(self) -> list[AssistantTool]:
         return tools
 
     async def run_async(self) -> "Run":
-        client = get_client()
+        client = get_openai_client()
 
         create_kwargs = {}
 
@@ -127,30 +127,33 @@ async def run_async(self) -> "Run":
         if self.tools is not None or self.additional_tools is not None:
             create_kwargs["tools"] = self.get_tools()
 
-        self.run = await client.beta.threads.runs.create(
-            thread_id=self.thread.id, assistant_id=self.assistant.id, **create_kwargs
-        )
-
-        self.assistant.pre_run_hook(run=self)
+        async with self.assistant:
+            self.run = await client.beta.threads.runs.create(
+                thread_id=self.thread.id,
+                assistant_id=self.assistant.id,
+                **create_kwargs,
+            )
 
-        try:
-            while self.run.status in ("queued", "in_progress", "requires_action"):
-                if self.run.status == "requires_action":
-                    await self._handle_step_requires_action()
-                await asyncio.sleep(0.1)
+            self.assistant.pre_run_hook(run=self)
+
+            try:
+                while self.run.status in ("queued", "in_progress", "requires_action"):
+                    if self.run.status == "requires_action":
+                        await self._handle_step_requires_action()
+                    await asyncio.sleep(0.1)
+                    await self.refresh_async()
+            except CancelRun as exc:
+                logger.debug(f"`CancelRun` raised; ending run with data: {exc.data}")
+                await client.beta.threads.runs.cancel(
+                    run_id=self.run.id, thread_id=self.thread.id
+                )
+                self.data = exc.data
                 await self.refresh_async()
-        except CancelRun as exc:
-            logger.debug(f"`CancelRun` raised; ending run with data: {exc.data}")
-            await client.beta.threads.runs.cancel(
-                run_id=self.run.id, thread_id=self.thread.id
-            )
-            self.data = exc.data
-            await self.refresh_async()
 
-        if self.run.status == "failed":
-            logger.debug(f"Run failed. Last error was: {self.run.last_error}")
+            if self.run.status == "failed":
+                logger.debug(f"Run failed. Last error was: {self.run.last_error}")
 
-        self.assistant.post_run_hook(run=self)
+            self.assistant.post_run_hook(run=self)
         return self
 
 
@@ -193,7 +196,7 @@ async def refresh_run_steps_async(self):
         max_attempts = max_fetched / limit + 2
 
         # Fetch the latest run steps
-        client = get_client()
+        client = get_openai_client()
 
         response = await client.beta.threads.runs.steps.list(
             run_id=self.run.id,