Merge pull request #5 from SpellcraftAI/dev

fix: large requests stalling, cleanup hanging

.github/workflows/release.yml

@@ -94,6 +94,7 @@ jobs:
 
   publish-to-testpypi:
     name: Publish Python 🐍 distribution 📦 to TestPyPI
+    if: contains(github.ref, 'refs/heads/dev')
     needs:
     - build
     runs-on: ubuntu-latest

oaib/Batch.py

@@ -344,11 +344,9 @@ async def run(self, callback=None):
         # existing requests first.
         if not self.__stopped.is_set():
             self.log("FINISHING PROCESSING | 5 second timeout")
-            await wait(
-                [*self.__processing, *self.__callbacks, *self.__workers],
-                return_when=ALL_COMPLETED,
-                timeout=5
-            )
+            await gather(*self.__processing)
+            await gather(*self.__workers)
+            await gather(*self.__callbacks)
             await self.stop()
 
         self.log("RETURNING OUTPUT")

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "oaib"
-version = "1.0.0"
+version = "1.0.1"
 requires-python = ">=3.9"
 description = "A Python library for making rate-limited, async batch requests to the OpenAI API."
 authors = [{ name = "CTJ Lewis", email = "[email protected]" }]

tests/test_large.py

@@ -0,0 +1,24 @@
+import oaib
+
+
+async def test_long():
+    batch = oaib.Auto()
+
+    n = 20
+    m = 20
+    for i in range(n):
+        await batch.add(
+            "chat.completions.create",
+            model="gpt-4",
+            max_tokens=4000,
+            messages=[{"role": "user", "content": "say hello and goodbye " * m}]
+        )
+
+    chats = await batch.run()
+    assert len(
+        chats) == n, f"Chat batch should return {n} results, got {len(chats)}"
+
+    chat = chats.iloc[0].get("result")
+    assert chat['choices'], "Should get valid chat completions"
+
+    chats