78f920eeda
Move common functions into their own folder and refactor the backends to use their own folder as well. Also cleanup imports and alphabetize import statments themselves. Finally, move colab and docker into their own folders as well. Signed-off-by: kingbri <bdashore3@proton.me>
27 lines
802 B
Python
27 lines
802 B
Python
"""Generator functions for the tabbyAPI."""
|
|
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):
|
|
"""Generate with a semaphore."""
|
|
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()
|