phoenix-oss/llama-stack
1
0
Fork 0
forked from phoenix-oss/llama-stack-mirror
Code Pull requests Releases Packages 2 Activity Actions

fix(library_client): ensure pending asyncio tasks like generator athrow are executed

Browse source
This commit is contained in:
Ashwin Bharambe 2025-03-09 16:17:27 -07:00
parent ba917a9c48
commit 70ff226b6a
1 changed files with 3 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

4
llama_stack/distribution/library_client.py
View file

@ -163,6 +163,9 @@ class LlamaStackAsLibraryClient(LlamaStackClient):
except StopAsyncIteration: except StopAsyncIteration:
pass pass
finally: finally:
pending = asyncio.all_tasks(loop)
if pending:
loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True))
loop.close() loop.close()
return sync_generator() return sync_generator()
@ -383,7 +386,6 @@ class AsyncLlamaStackAsLibraryClient(AsyncLlamaStackClient):
# Wrap the generator to preserve context across iterations # Wrap the generator to preserve context across iterations
wrapped_gen = preserve_headers_context_async_generator(gen()) wrapped_gen = preserve_headers_context_async_generator(gen())
mock_response = httpx.Response( mock_response = httpx.Response(
status_code=httpx.codes.OK, status_code=httpx.codes.OK,
content=wrapped_gen, content=wrapped_gen,