| 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""" |
| |
| caller_frame = inspect.stack()[1] |
| caller_filepath = caller_frame.filename |
|
|
| |
| caller_dir = os.path.dirname(os.path.abspath(caller_filepath)) |
|
|
| return os.path.join(caller_dir, path) |
|
|
