Archives/langchain
2
0
Fork 1
mirror of https://github.com/hwchase17/langchain synced 2024-11-06 03:20:49 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix async task being destroyed before cancelled (#2787)

Browse Source
This commit is contained in:
Nuno Campos 2023-04-12 20:38:38 +01:00 committed by GitHub
parent 819d72614a
commit 1bea9ea4be
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 6 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
langchain/callbacks/streaming_aiter.py
View File

@ -47,14 +47,19 @@ class AsyncIteratorCallbackHandler(AsyncCallbackHandler):
while not self.queue.empty() or not self.done.is_set(): while not self.queue.empty() or not self.done.is_set():
# Wait for the next token in the queue, # Wait for the next token in the queue,
# but stop waiting if the done event is set # but stop waiting if the done event is set
done, _ = await asyncio.wait( done, other = await asyncio.wait(
[ [
# NOTE: If you add other tasks here, update the code below,
# which assumes each set has exactly one task each
asyncio.ensure_future(self.queue.get()), asyncio.ensure_future(self.queue.get()),
asyncio.ensure_future(self.done.wait()), asyncio.ensure_future(self.done.wait()),
], ],
return_when=asyncio.FIRST_COMPLETED, return_when=asyncio.FIRST_COMPLETED,
) )
# Cancel the other task
other.pop().cancel()
# Extract the value of the first completed task # Extract the value of the first completed task
token_or_done = cast(Union[str, Literal[True]], done.pop().result()) token_or_done = cast(Union[str, Literal[True]], done.pop().result())