# -*- coding: utf-8 -*-
"""A manager for AgentScope."""
import os
from typing import Union, Any
from copy import deepcopy
from loguru import logger
from ._monitor import MonitorManager
from ._file import FileManager
from ._model import ModelManager
from ..logging import LOG_LEVEL, setup_logger
from .._version import __version__
from ..utils.common import (
_generate_random_code,
_get_process_creation_time,
_get_timestamp,
)
from ..constants import _RUNTIME_ID_FORMAT, _RUNTIME_TIMESTAMP_FORMAT
from ..studio._client import _studio_client
[docs]
class ASManager:
"""A manager for AgentScope."""
_instance = None
__serialized_attrs = [
"project",
"name",
"disable_saving",
"run_id",
"pid",
"timestamp",
]
def __new__(cls, *args: Any, **kwargs: Any) -> "ASManager":
if cls._instance is None:
cls._instance = super(ASManager, cls).__new__(cls)
return cls._instance
[docs]
@classmethod
def get_instance(cls) -> "ASManager":
"""Get the instance of the singleton class."""
if cls._instance is None:
raise ValueError(
"AgentScope hasn't been initialized. Please call "
"`agentscope.init` function first.",
)
return cls._instance
[docs]
def __init__(self) -> None:
"""Initialize the manager. Note we initialize the managers by default
arguments to avoid unnecessary errors when user doesn't call
`agentscope.init` function"""
self.project = ""
self.name = ""
self.run_id = ""
self.pid = -1
self.timestamp = ""
self.disable_saving = True
self.file = FileManager()
self.model = ModelManager()
self.monitor = MonitorManager()
# TODO: unified with logger and studio
self.logger_level: LOG_LEVEL = "INFO"
[docs]
def initialize(
self,
model_configs: Union[dict, str, list, None],
project: Union[str, None],
name: Union[str, None],
disable_saving: bool,
save_dir: str,
save_log: bool,
save_code: bool,
save_api_invoke: bool,
cache_dir: str,
use_monitor: bool,
logger_level: LOG_LEVEL,
run_id: Union[str, None],
studio_url: Union[str, None],
) -> None:
"""Initialize the package."""
# =============== Init the runtime ===============
self.project = project or _generate_random_code()
self.name = name or _generate_random_code(uppercase=False)
self.pid = os.getpid()
timestamp = _get_process_creation_time()
self.timestamp = timestamp.strftime(_RUNTIME_TIMESTAMP_FORMAT)
self.run_id = run_id or _get_timestamp(
_RUNTIME_ID_FORMAT,
timestamp,
).format(self.name)
self.disable_saving = disable_saving
# =============== Init the file manager ===============
if disable_saving:
save_log = False
save_code = False
save_api_invoke = False
use_monitor = False
run_dir = None
else:
run_dir = os.path.abspath(os.path.join(save_dir, self.run_id))
self.file.initialize(
run_dir=run_dir,
save_log=save_log,
save_code=save_code,
save_api_invoke=save_api_invoke,
cache_dir=cache_dir,
)
# Save the python code here to avoid duplicated saving in the child
# process (when calling deserialize function)
if save_code:
self.file.save_python_code()
if not disable_saving:
# Save the runtime information in .config file
self.file.save_runtime_information(self.state_dict())
# =============== Init the logger ===============
# TODO: unified with studio and gradio
self.logger_level = logger_level
# run_dir will be None if save_log is False
setup_logger(self.file.run_dir, logger_level)
# =============== Init the model manager ===============
self.model.initialize(model_configs)
# =============== Init the monitor manager ===============
self.monitor.initialize(use_monitor)
# =============== Init the studio ===============
# TODO: unified with studio and gradio
# Init studio client, which will push messages to web ui and fetch user
# inputs from web ui
if studio_url is not None:
_studio_client.initialize(self.run_id, studio_url)
# Register in AgentScope Studio
_studio_client.register_running_instance(
project=self.project,
name=self.name,
timestamp=self.timestamp,
run_dir=self.file.run_dir,
pid=self.pid,
)
[docs]
def state_dict(self) -> dict:
"""Serialize the runtime information."""
serialized_data = {
k: getattr(self, k) for k in self.__serialized_attrs
}
serialized_data["agentscope_version"] = __version__
serialized_data["file"] = self.file.state_dict()
serialized_data["model"] = self.model.state_dict()
serialized_data["logger"] = {
"level": self.logger_level,
}
serialized_data["studio"] = _studio_client.state_dict()
serialized_data["monitor"] = self.monitor.state_dict()
return deepcopy(serialized_data)
[docs]
def load_dict(self, data: dict) -> None:
"""Load the runtime information from a dictionary"""
for k in self.__serialized_attrs:
assert k in data, f"Key {k} not found in data."
setattr(self, k, data[k])
self.file.load_dict(data["file"])
# TODO: unified the logger with studio and gradio
self.logger_level = data["logger"]["level"]
setup_logger(self.file.run_dir, self.logger_level)
self.model.load_dict(data["model"])
_studio_client.load_dict(data["studio"])
self.monitor.load_dict(data["monitor"])
[docs]
def flush(self) -> None:
"""Flush the runtime information."""
self.project = ""
self.name = ""
self.run_id = ""
self.pid = -1
self.timestamp = ""
self.disable_saving = True
self.file.flush()
self.model.flush()
self.monitor.flush()
logger.remove()
_studio_client.flush()
self.logger_level = "INFO"