# -*- coding: utf-8 -*-
""" Common utils."""
import base64
import contextlib
import datetime
import hashlib
import json
import os
import random
import re
import secrets
import signal
import socket
import string
import sys
import tempfile
import threading
from typing import Any, Generator, Optional, Union, Tuple, Literal, List
from urllib.parse import urlparse
import psutil
import requests
[文档]
@contextlib.contextmanager
def timer(seconds: Optional[Union[int, float]] = None) -> Generator:
"""
A context manager that limits the execution time of a code block to a
given number of seconds.
The implementation of this contextmanager are borrowed from
https://github.com/openai/human-eval/blob/master/human_eval/execution.py
Note:
This function only works in Unix and MainThread,
since `signal.setitimer` is only available in Unix.
"""
if (
seconds is None
or sys.platform == "win32"
or threading.currentThread().name # pylint: disable=W4902
!= "MainThread"
):
yield
return
def signal_handler(*args: Any, **kwargs: Any) -> None:
raise TimeoutError("timed out")
signal.setitimer(signal.ITIMER_REAL, seconds)
signal.signal(signal.SIGALRM, signal_handler)
try:
# Enter the context and execute the code block.
yield
finally:
signal.setitimer(signal.ITIMER_REAL, 0)
[文档]
@contextlib.contextmanager
def create_tempdir() -> Generator:
"""
A context manager that creates a temporary directory and changes the
current working directory to it.
The implementation of this contextmanager are borrowed from
https://github.com/openai/human-eval/blob/master/human_eval/execution.py
"""
with tempfile.TemporaryDirectory() as dirname:
with _chdir(dirname):
yield dirname
@contextlib.contextmanager
def _chdir(path: str) -> Generator:
"""
A context manager that changes the current working directory to the
given path.
The implementation of this contextmanager are borrowed from
https://github.com/openai/human-eval/blob/master/human_eval/execution.py
"""
if path == ".":
yield
return
cwd = os.getcwd()
os.chdir(path)
try:
yield
except BaseException as exc:
raise exc
finally:
os.chdir(cwd)
def _requests_get(
url: str,
params: dict,
headers: Optional[dict] = None,
) -> Union[dict, str]:
"""
Sends a GET request to the specified URL with the provided query parameters
and headers. Returns the JSON response as a dictionary.
This function handles the request, checks for errors, logs exceptions,
and parses the JSON response.
Args:
url (str): The URL to which the GET request is sent.
params (Dict): A dictionary containing query parameters to be included
in the request.
headers (Optional[Dict]): An optional dictionary of HTTP headers to
send with the request.
Returns:
Dict or str:
If the request is successful, returns a dictionary containing the
parsed JSON data.
If the request fails, returns the error string.
"""
# Make the request
try:
# Check if headers are provided, and include them if they are not None
if headers:
response = requests.get(url, params=params, headers=headers)
else:
response = requests.get(url, params=params)
# This will raise an exception for HTTP error codes
response.raise_for_status()
except requests.RequestException as e:
return str(e)
# Parse the JSON response
search_results = response.json()
return search_results
def _if_change_database(sql_query: str) -> bool:
"""Check whether SQL query only contains SELECT query"""
# Compile the regex pattern outside the function for better performance
pattern_unsafe_sql = re.compile(
r"\b(INSERT|UPDATE|DELETE|REPLACE|CREATE|ALTER|DROP|TRUNCATE|USE)\b",
re.IGNORECASE,
)
# Remove SQL comments
sql_query = re.sub(r"--.*?$", "", sql_query, flags=re.MULTILINE)
# Remove /* */ comments
sql_query = re.sub(r"/\*.*?\*/", "", sql_query, flags=re.DOTALL)
# Matching non-SELECT statements with regular expressions
if pattern_unsafe_sql.search(sql_query):
return False
return True
def _get_timestamp(
format_: str = "%Y-%m-%d %H:%M:%S",
time: datetime.datetime = None,
) -> str:
"""Get current timestamp."""
if time is None:
return datetime.datetime.now().strftime(format_)
else:
return time.strftime(format_)
[文档]
def to_openai_dict(item: dict) -> dict:
"""Convert `Msg` to `dict` for OpenAI API."""
clean_dict = {}
if "name" in item:
clean_dict["name"] = item["name"]
if "role" in item:
clean_dict["role"] = item["role"]
else:
clean_dict["role"] = "assistant"
if "content" in item:
clean_dict["content"] = _convert_to_str(item["content"])
else:
raise ValueError("The content of the message is missing.")
return clean_dict
def _find_available_port() -> int:
"""Get an unoccupied socket port number."""
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.bind(("", 0))
return s.getsockname()[1]
def _check_port(port: Optional[int] = None) -> int:
"""Check if the port is available.
Args:
port (`int`):
the port number being checked.
Returns:
`int`: the port number that passed the check. If the port is found
to be occupied, an available port number will be automatically
returned.
"""
if port is None:
new_port = _find_available_port()
return new_port
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
try:
if s.connect_ex(("localhost", port)) == 0:
raise RuntimeError("Port is occupied.")
except Exception:
new_port = _find_available_port()
return new_port
return port
def _guess_type_by_extension(
url: str,
) -> Literal["image", "audio", "video", "file"]:
"""Guess the type of the file by its extension."""
extension = url.split(".")[-1].lower()
if extension in [
"bmp",
"dib",
"icns",
"ico",
"jfif",
"jpe",
"jpeg",
"jpg",
"j2c",
"j2k",
"jp2",
"jpc",
"jpf",
"jpx",
"apng",
"png",
"bw",
"rgb",
"rgba",
"sgi",
"tif",
"tiff",
"webp",
]:
return "image"
elif extension in [
"amr",
"wav",
"3gp",
"3gpp",
"aac",
"mp3",
"flac",
"ogg",
]:
return "audio"
elif extension in [
"mp4",
"webm",
"mkv",
"flv",
"avi",
"mov",
"wmv",
"rmvb",
]:
return "video"
else:
return "file"
def _to_openai_image_url(url: str) -> str:
"""Convert an image url to openai format. If the given url is a local
file, it will be converted to base64 format. Otherwise, it will be
returned directly.
Args:
url (`str`):
The local or public url of the image.
"""
# See https://platform.openai.com/docs/guides/vision for details of
# support image extensions.
support_image_extensions = (
".png",
".jpg",
".jpeg",
".gif",
".webp",
)
parsed_url = urlparse(url)
lower_url = url.lower()
# Web url
if parsed_url.scheme != "":
if any(lower_url.endswith(_) for _ in support_image_extensions):
return url
# Check if it is a local file
elif os.path.exists(url) and os.path.isfile(url):
if any(lower_url.endswith(_) for _ in support_image_extensions):
with open(url, "rb") as image_file:
base64_image = base64.b64encode(image_file.read()).decode(
"utf-8",
)
extension = parsed_url.path.lower().split(".")[-1]
mime_type = f"image/{extension}"
return f"data:{mime_type};base64,{base64_image}"
raise TypeError(f"{url} should be end with {support_image_extensions}.")
def _download_file(url: str, path_file: str, max_retries: int = 3) -> bool:
"""Download file from the given url and save it to the given path.
Args:
url (`str`):
The url of the file.
path_file (`str`):
The path to save the file.
max_retries (`int`, defaults to `3`)
The maximum number of retries when fail to download the file.
"""
for n_retry in range(1, max_retries + 1):
response = requests.get(url, stream=True)
if response.status_code == requests.codes.ok:
with open(path_file, "wb") as file:
for chunk in response.iter_content(1024):
file.write(chunk)
return True
else:
raise RuntimeError(
f"Failed to download file from {url} (status code: "
f"{response.status_code}). Retry {n_retry}/{max_retries}.",
)
return False
def _generate_random_code(
length: int = 6,
uppercase: bool = True,
lowercase: bool = True,
digits: bool = True,
) -> str:
"""Get random code."""
characters = ""
if uppercase:
characters += string.ascii_uppercase
if lowercase:
characters += string.ascii_lowercase
if digits:
characters += string.digits
return "".join(secrets.choice(characters) for i in range(length))
def _generate_id_from_seed(seed: str, length: int = 8) -> str:
"""Generate random id from seed str.
Args:
seed (`str`): seed string.
length (`int`): generated id length.
"""
hasher = hashlib.sha256()
hasher.update(seed.encode("utf-8"))
hash_digest = hasher.hexdigest()
random.seed(hash_digest)
id_chars = [
random.choice(string.ascii_letters + string.digits)
for _ in range(length)
]
return "".join(id_chars)
def _is_web_url(url: str) -> bool:
"""Whether the url is accessible from the Web.
Args:
url (`str`):
The url to check.
Note:
This function is not perfect, it only checks if the URL starts with
common web protocols, e.g., http, https, ftp, oss.
"""
parsed_url = urlparse(url)
return parsed_url.scheme in ["http", "https", "ftp", "oss"]
def _is_json_serializable(obj: Any) -> bool:
"""Check if the given object is json serializable."""
try:
json.dumps(obj)
return True
except TypeError:
return False
def _convert_to_str(content: Any) -> str:
"""Convert the content to string.
Note:
For prompt engineering, simply calling `str(content)` or
`json.dumps(content)` is not enough.
- For `str(content)`, if `content` is a dictionary, it will turn double
quotes to single quotes. When this string is fed into prompt, the LLMs
may learn to use single quotes instead of double quotes (which
cannot be loaded by `json.loads` API).
- For `json.dumps(content)`, if `content` is a string, it will add
double quotes to the string. LLMs may learn to use double quotes to
wrap strings, which leads to the same issue as `str(content)`.
To avoid these issues, we use this function to safely convert the
content to a string used in prompt.
Args:
content (`Any`):
The content to be converted.
Returns:
`str`: The converted string.
"""
if isinstance(content, str):
return content
elif isinstance(content, (dict, list, int, float, bool, tuple)):
return json.dumps(content, ensure_ascii=False)
else:
return str(content)
def _join_str_with_comma_and(elements: List[str]) -> str:
"""Return the JSON string with comma, and use " and " between the last two
elements."""
if len(elements) == 0:
return ""
elif len(elements) == 1:
return elements[0]
elif len(elements) == 2:
return " and ".join(elements)
else:
return ", ".join(elements[:-1]) + f", and {elements[-1]}"
[文档]
class ImportErrorReporter:
"""Used as a placeholder for missing packages.
When called, an ImportError will be raised, prompting the user to install
the specified extras requirement.
"""
[文档]
def __init__(self, error: ImportError, extras_require: str = None) -> None:
"""Init the ImportErrorReporter.
Args:
error (`ImportError`): the original ImportError.
extras_require (`str`): the extras requirement.
"""
self.error = error
self.extras_require = extras_require
def __call__(self, *args: Any, **kwds: Any) -> Any:
return self._raise_import_error()
def __getattr__(self, name: str) -> Any:
return self._raise_import_error()
def __getitem__(self, __key: Any) -> Any:
return self._raise_import_error()
def _raise_import_error(self) -> Any:
"""Raise the ImportError"""
err_msg = f"ImportError occorred: [{self.error.msg}]."
if self.extras_require is not None:
err_msg += (
f" Please install [{self.extras_require}] version"
" of agentscope."
)
raise ImportError(err_msg)
def _hash_string(
data: str,
hash_method: Literal["sha256", "md5", "sha1"],
) -> str:
"""Hash the string data."""
hash_func = getattr(hashlib, hash_method)()
hash_func.update(data.encode())
return hash_func.hexdigest()
def _get_process_creation_time() -> datetime.datetime:
"""Get the creation time of the process."""
pid = os.getpid()
# Find the process by pid
current_process = psutil.Process(pid)
# Obtain the process creation time
create_time = current_process.create_time()
# Change the timestamp to a readable format
return datetime.datetime.fromtimestamp(create_time)
def _is_process_alive(
pid: int,
create_time_str: str,
create_time_format: str = "%Y-%m-%d %H:%M:%S",
tolerance_seconds: int = 10,
) -> bool:
"""Check if the process is alive by comparing the actual creation time of
the process with the given creation time.
Args:
pid (`int`):
The process id.
create_time_str (`str`):
The given creation time string.
create_time_format (`str`, defaults to `"%Y-%m-%d %H:%M:%S"`):
The format of the given creation time string.
tolerance_seconds (`int`, defaults to `10`):
The tolerance seconds for comparing the actual creation time with
the given creation time.
Returns:
`bool`: True if the process is alive, False otherwise.
"""
try:
# Try to create a process object by pid
proc = psutil.Process(pid)
# Obtain the actual creation time of the process
actual_create_time_timestamp = proc.create_time()
# Convert the given creation time string to a datetime object
given_create_time_datetime = datetime.datetime.strptime(
create_time_str,
create_time_format,
)
# Calculate the time difference between the actual creation time and
time_difference = abs(
actual_create_time_timestamp
- given_create_time_datetime.timestamp(),
)
# Compare the actual creation time with the given creation time
if time_difference <= tolerance_seconds:
return True
except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
# If the process is not found, access is denied, or the process is a
# zombie process, return False
return False
return False
def _is_windows() -> bool:
"""Check if the system is Windows."""
return os.name == "nt"
def _map_string_to_color_mark(
target_str: str,
) -> Tuple[str, str]:
"""Map a string into an index within a given length.
Args:
target_str (`str`):
The string to be mapped.
Returns:
`Tuple[str, str]`: A color marker tuple
"""
color_marks = [
("\033[90m", "\033[0m"),
("\033[91m", "\033[0m"),
("\033[92m", "\033[0m"),
("\033[93m", "\033[0m"),
("\033[94m", "\033[0m"),
("\033[95m", "\033[0m"),
("\033[96m", "\033[0m"),
("\033[97m", "\033[0m"),
]
hash_value = int(hashlib.sha256(target_str.encode()).hexdigest(), 16)
index = hash_value % len(color_marks)
return color_marks[index]
def _generate_new_runtime_id() -> str:
"""Generate a new random runtime id."""
_RUNTIME_ID_FORMAT = "run_%Y%m%d-%H%M%S_{}"
return _get_timestamp(_RUNTIME_ID_FORMAT).format(
_generate_random_code(uppercase=False),
)