mirror of
https://github.com/theroyallab/tabbyAPI.git
synced 2026-03-14 15:57:27 +00:00
1a8afcb6ad
Non-streaming tasks were not regulated by the semaphore, causing these tasks to interfere with streaming generations. Add helper functions to take in both sync and async functions for callbacks and sequential blocking with the semaphore. Signed-off-by: kingbri <bdashore3@proton.me>
24 lines
719 B
Python
24 lines
719 B
Python
import inspect
|
|
from asyncio import Semaphore
|
|
from functools import partialmethod
|
|
from typing import AsyncGenerator
|
|
|
|
generate_semaphore = Semaphore(1)
|
|
|
|
# Async generation that blocks on a semaphore
|
|
async def generate_with_semaphore(generator: AsyncGenerator):
|
|
async with generate_semaphore:
|
|
if inspect.isasyncgenfunction:
|
|
async for result in generator():
|
|
yield result
|
|
else:
|
|
for result in generator():
|
|
yield result
|
|
|
|
# Block a function with semaphore
|
|
async def call_with_semaphore(callback: partialmethod):
|
|
if inspect.iscoroutinefunction(callback):
|
|
return await callback()
|
|
async with generate_semaphore:
|
|
return callback()
|