From fbf127d7863bbe1d4250af5657c2641700da9ba0 Mon Sep 17 00:00:00 2001
From: Nick Hale <4175918+njhale@users.noreply.github.com>
Date: Mon, 8 Jul 2024 11:16:05 -0400
Subject: [PATCH] fix: gracefully exit on interrupt

Catch KeyboardInterrupt and asyncio.CancelledError exceptions to ensure
the application exits without printing traceback on interrupt. This
allows gptscript to kill the provider silently.

Signed-off-by: Nick Hale <4175918+njhale@users.noreply.github.com>
---
 main.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/main.py b/main.py
index 6d167d8..6c031e8 100644
--- a/main.py
+++ b/main.py
@@ -111,6 +111,10 @@ async def convert_stream(stream: Stream[ChatCompletionChunk]) -> AsyncIterable[s
 
 if __name__ == "__main__":
     import uvicorn
+    import asyncio
 
-    uvicorn.run("main:app", host="127.0.0.1", port=int(os.environ.get("PORT", "8000")),
+    try:
+        uvicorn.run("main:app", host="127.0.0.1", port=int(os.environ.get("PORT", "8000")),
                 log_level="debug" if debug else "critical", reload=debug, access_log=debug)
+    except (KeyboardInterrupt, asyncio.CancelledError):
+        pass