| import asyncio |
| import importlib |
| import inspect |
| import logging |
| import os |
| import os.path as osp |
| import sys |
| import time |
| from functools import partial |
| from logging.handlers import RotatingFileHandler |
| from typing import Any, Dict, Generator, Iterable, List, Optional, Union |
|
|
|
|
| def load_class_from_string(class_path: str, path=None): |
| path_in_sys = False |
| if path: |
| if path not in sys.path: |
| path_in_sys = True |
| sys.path.insert(0, path) |
|
|
| try: |
| module_name, class_name = class_path.rsplit('.', 1) |
| module = importlib.import_module(module_name) |
| cls = getattr(module, class_name) |
| return cls |
| finally: |
| if path and path_in_sys: |
| sys.path.remove(path) |
|
|
|
|
| def create_object(config: Union[Dict, Any] = None): |
| """Create an instance based on the configuration where 'type' is a |
| preserved key to indicate the class (path). When accepting non-dictionary |
| input, the function degenerates to an identity. |
| """ |
| if config is None or not isinstance(config, dict): |
| return config |
| assert isinstance(config, dict) and 'type' in config |
|
|
| config = config.copy() |
| obj_type = config.pop('type') |
| if isinstance(obj_type, str): |
| obj_type = load_class_from_string(obj_type) |
| if inspect.isclass(obj_type): |
| obj = obj_type(**config) |
| else: |
| assert callable(obj_type) |
| obj = partial(obj_type, **config) |
| return obj |
|
|
|
|
| async def async_as_completed(futures: Iterable[asyncio.Future]): |
| """A asynchronous wrapper for `asyncio.as_completed`""" |
| loop = asyncio.get_event_loop() |
| wrappers = [] |
| for fut in futures: |
| assert isinstance(fut, asyncio.Future) |
| wrapper = loop.create_future() |
| fut.add_done_callback(wrapper.set_result) |
| wrappers.append(wrapper) |
| for next_completed in asyncio.as_completed(wrappers): |
| yield await next_completed |
|
|
|
|
| def filter_suffix(response: Union[str, List[str]], |
| suffixes: Optional[List[str]] = None) -> str: |
| """Filter response with suffixes. |
| |
| Args: |
| response (Union[str, List[str]]): generated responses by LLMs. |
| suffixes (str): a list of suffixes to be deleted. |
| |
| Return: |
| str: a clean response. |
| """ |
| if suffixes is None: |
| return response |
| batched = True |
| if isinstance(response, str): |
| response = [response] |
| batched = False |
| processed = [] |
| for resp in response: |
| for item in suffixes: |
| |
| |
| if item in resp: |
| resp = resp.split(item)[0] |
| processed.append(resp) |
| if not batched: |
| return processed[0] |
| return processed |
|
|
|
|
| def get_logger( |
| name: str = 'lagent', |
| level: str = 'debug', |
| fmt: |
| str = '%(asctime)s %(levelname)8s %(filename)20s %(lineno)4s - %(message)s', |
| add_file_handler: bool = False, |
| log_dir: str = 'log', |
| log_file: str = time.strftime('%Y-%m-%d.log', time.localtime()), |
| max_bytes: int = 5 * 1024 * 1024, |
| backup_count: int = 3, |
| ): |
| logger = logging.getLogger(name) |
| logger.propagate = False |
| logger.setLevel(getattr(logging, level.upper(), logging.DEBUG)) |
|
|
| formatter = logging.Formatter(fmt) |
| console_handler = logging.StreamHandler() |
| console_handler.setFormatter(formatter) |
| logger.addHandler(console_handler) |
|
|
| if add_file_handler: |
| if not osp.exists(log_dir): |
| os.makedirs(log_dir) |
| log_file_path = osp.join(log_dir, log_file) |
| file_handler = RotatingFileHandler( |
| log_file_path, |
| maxBytes=max_bytes, |
| backupCount=backup_count, |
| encoding='utf-8') |
| file_handler.setFormatter(formatter) |
| logger.addHandler(file_handler) |
|
|
| return logger |
|
|
|
|
| class GeneratorWithReturn: |
| """Generator wrapper to capture the return value.""" |
|
|
| def __init__(self, generator: Generator): |
| self.generator = generator |
| self.ret = None |
|
|
| def __iter__(self): |
| self.ret = yield from self.generator |
| return self.ret |
|
|
