Merge pull request #25 from thedadams/await-event-handlers

fix: await event_handlers to ensure errors are processed
gptscript-ai · Jun 24, 2024 · d85b623 · d85b623
2 parents 3e8cd9e + 5b9473d
commit d85b623
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 3 deletions.
diff --git a/gptscript/run.py b/gptscript/run.py
@@ -33,6 +33,7 @@ def __init__(self, subCommand: str, tools: Union[ToolDef | list[ToolDef] | str],
         self._rawOutput: Any = None
         self._task: Awaitable | None = None
         self._resp: httpx.Response | None = None
+        self._event_tasks: list[Awaitable[None]] = []
 
     def program(self):
         return self._program
@@ -157,7 +158,7 @@ async def _request(self, tool: Any):
                                 self._parentCallID = event.id
                         if self.event_handlers is not None:
                             for event_handler in self.event_handlers:
-                                asyncio.create_task(event_handler(self, event))
+                                self._event_tasks.append(asyncio.create_task(event_handler(self, event)))
 
         self._resp = None
         if self._err != "":
@@ -167,6 +168,14 @@ async def _request(self, tool: Any):
         else:
             self._state = RunState.Continue
 
+        for task in self._event_tasks:
+            try:
+                await task
+            except Exception as e:
+                print(f"error during event processing: {e}")
+
+        self._event_tasks = []
+
     async def aclose(self):
         if self._task is None or self._resp is None:
             raise Exception("run not started")

diff --git a/tests/test_gptscript.py b/tests/test_gptscript.py
@@ -22,7 +22,8 @@ def gptscript():
         pytest.fail("OPENAI_API_KEY not set", pytrace=False)
     try:
         gptscript = GPTScript(GlobalOptions(apiKey=os.getenv("OPENAI_API_KEY")))
-        return gptscript
+        yield gptscript
+        gptscript.close()
     except Exception as e:
         pytest.fail(e, pytrace=False)
 
@@ -387,7 +388,7 @@ async def process_event(r: Run, frame: CallFrame | RunFrame | PromptFrame):
         nonlocal confirm_event_found, event_content
         if frame.type == RunEventType.callConfirm:
             confirm_event_found = True
-            assert '"ls"' in frame.input or '"dir"' in frame.input, "Unexpected confirm input: " + frame.input
+            assert '"ls' in frame.input or '"dir' in frame.input, "Unexpected confirm input: " + frame.input
             await gptscript.confirm(AuthResponse(frame.id, True))
         elif frame.type == RunEventType.callProgress:
             for output in frame.output: