File size: 1,990 Bytes
42da440 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 | import asyncio
from augmentoolkit.generation_functions.engine_wrapper_class import EngineWrapper
import os
import inspect
def setup_semaphore_and_engines(
concurrency_limit: int,
small_model: str,
small_api_key: str,
small_base_url: str,
small_mode: str,
large_model: str,
large_api_key: str,
large_base_url: str,
large_mode: str,
engine_input_observers=[],
engine_output_observers=[],
large_engine_input_observers=[],
large_engine_output_observers=[],
):
semaphore = asyncio.Semaphore(concurrency_limit)
async def run_task_with_limit(task, timeout=60):
async with semaphore:
try:
return await asyncio.wait_for(task, timeout=timeout)
except asyncio.TimeoutError:
print("[Timeout] A task exceeded the timeout limit.")
return None
except Exception as e:
print(f"[Error] Exception in task: {e}")
return None
engine_wrapper = EngineWrapper(
model=small_model,
api_key=small_api_key,
base_url=small_base_url,
mode=small_mode,
input_observers=engine_input_observers,
output_observers=engine_output_observers,
)
engine_wrapper_large = EngineWrapper(
model=large_model,
api_key=large_api_key,
base_url=large_base_url,
mode=large_mode,
input_observers=large_engine_input_observers,
output_observers=large_engine_output_observers,
)
return run_task_with_limit, engine_wrapper, engine_wrapper_large, semaphore
def make_relative_to_self(path):
"""Make paths relative to the calling script's location"""
# Get the caller's frame info
caller_frame = inspect.stack()[1]
caller_filepath = caller_frame.filename
# Get directory of the script that called this function
caller_dir = os.path.dirname(os.path.abspath(caller_filepath))
return os.path.join(caller_dir, path)
|