File size: 12,796 Bytes
77320e4 |
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 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 |
import asyncio
from abc import ABC, abstractmethod
from typing import Dict, Callable, Union, AsyncGenerator
from ..exceptions.exceptions import InputErrorException
from ..prompt import PromptTemplate
from ..schemas import AgentOutput, AgentType, AgentResponse
from ..llm.base_llm import BaseLLM
from ..tools import BaseTool
from ..utils import Config, get_logger
import os
from importlib import import_module
logger = get_logger()
LLM_CONF_OVERRIDE_KEY = ['psm', 'dc', 'temperature', 'top_p', 'top_k', 'max_tokens']
class BaseAgent(ABC):
"""Base Agent class defining the essential attributes and methods for an ALM Agent.
"""
def __init__(self, **kwargs):
"""
Initializes an instance of the Agent class.
"""
# Set default values
default_config = {
'name': 'agent',
'type': AgentType.react,
'version': '',
'description': '',
'prompt_template': None,
'auth': {}
}
# Update default values with provided config
default_config.update(kwargs)
# Access configuration data with a known default value
auth = default_config['auth']
self._set_auth_env(auth)
self._name: str = default_config['name']
self._type: AgentType = default_config['type']
self._version: str = default_config['version']
self._description: str = default_config['description']
self.__prompt_template: Union[PromptTemplate, None] = \
self._get_prompt_template(default_config['prompt_template'])
self.__llm: Union[BaseLLM, None] = None
self.__plugins_map: Dict = {}
self.__plugin_tool_function = {}
self.__plugin_tool_async_function = {}
self.__plugin_tool_description = None
@property
def name(self) -> str:
return self._name
@property
def type(self) -> AgentType:
return self._type
@property
def version(self) -> str:
return self._version
@property
def description(self) -> str:
return self._description
@property
def prompt_template(self) -> PromptTemplate:
return self.__prompt_template
@property
def llm(self) -> Union[BaseLLM, None]:
return self.__llm
@llm.setter
def llm(self, llm_client: BaseLLM):
if llm_client is None or not isinstance(llm_client, BaseLLM):
raise InputErrorException("Invalid llm client {}".format(type(llm_client)))
self.__llm = llm_client
@property
def plugins_map(self) -> Dict:
return self.__plugins_map.copy() # Return a copy to prevent external modification
def add_plugin(self, tool_name: str, tool):
if not tool_name or not tool:
raise InputErrorException("Adding invalid tool name: {}, type {}".format(tool_name, type(tool)))
self.__plugins_map[tool_name] = tool
def _set_auth_env(self, obj):
"""This method sets environment variables for authentication.
"""
for key in obj:
os.environ[key] = obj.get(key)
def _get_prompt_template(self, obj):
"""This method returns a prompt template instance based on the provided configuration.
"""
assert isinstance(obj, dict) or isinstance(obj, PromptTemplate)
if isinstance(obj, dict):
return {
key: self._parse_prompt_template(obj[key]) for key in obj
}
elif isinstance(obj, PromptTemplate):
ans = self._parse_prompt_template(obj)
return ans
else:
raise InputErrorException("Invalid PromptTemplate, it should be a dict or PromptTemplate. But get {}"
.format(type(obj)))
def _parse_prompt_template(self, obj: Union[dict, PromptTemplate]):
"""This method parses the prompt template configuration and returns a prompt template instance.
"""
assert isinstance(obj, dict) or isinstance(obj, PromptTemplate)
if isinstance(obj, PromptTemplate):
return obj
return PromptTemplate(input_variables=obj['input_variables'],
template=obj['template'],
validate_template=bool(obj.get('validate_template', True)))
@classmethod
def _get_basic_instance_from_config(cls, config_data):
agent_module_name = config_data.get("module_name", None)
agent_class_name = config_data.get("class_name", None)
if not agent_module_name or not agent_class_name:
raise InputErrorException("Agent module_name and class_name required, please check your config")
module = import_module(agent_module_name)
clazz = getattr(module, agent_class_name)
agent_instance = clazz(**config_data)
return agent_instance
@classmethod
def from_config_path_and_kwargs(cls, config_path, **kwargs):
config_data = Config.load(config_path)
logger.info(f"Use config from path {config_path} to init agent : {config_data}")
agent_instance = cls._get_basic_instance_from_config(config_data)
if 'llm' in config_data and 'params' in config_data['llm']:
for param in LLM_CONF_OVERRIDE_KEY:
if param in kwargs and kwargs[param]:
logger.info(f"Overwrite with new {param} {kwargs[param]}")
config_data['llm']['params'][param] = kwargs[param]
assert isinstance(agent_instance, BaseAgent)
agent_instance._init_llm(config_data.get("llm", {}))
agent_instance._init_plugins(config_data.get('plugins', []))
return agent_instance
def _init_llm(self, obj):
"""
This method parses the Language Model Manager (LLM) configuration and returns an LLM instance.
:param obj: A configuration dictionary or string.
:type obj: dict or str
:raises ValueError: If the specified LLM is not supported.
:return: An LLM instance.
:rtype: BaseLLM
"""
if isinstance(obj, str):
name = obj
model_params = dict()
else:
name = obj.get('model_name', None)
model_params = obj.get('params', dict())
module_name = obj['module_name']
class_name = obj['class_name']
module = import_module(module_name)
clazz = getattr(module, class_name)
llm = clazz(model_name=name, params=model_params)
self.llm = llm
def _init_plugins(self, configs):
"""
This method parses the plugin configuration and add each plugin into the plugins_map.
"""
assert isinstance(configs, list)
for plugin_config in configs:
if plugin_config.get('type', "") == 'agent':
# Agent as plugin
agent = BaseAgent.from_config_path_and_kwargs(plugin_config['config'])
self.plugins_map[plugin_config['name']] = agent
else:
# Tools as plugin
params = plugin_config.get('params', dict())
tool = BaseTool.from_config(config_input=plugin_config['config'], **params)
self.plugins_map[tool.name] = tool
@classmethod
async def async_from_config_path_and_kwargs(cls, config_path, **kwargs):
config_data = Config.load(config_path)
logger.info(f"Use config from path {config_path} to init agent : {config_data}")
agent_instance = cls._get_basic_instance_from_config(config_data)
# override default config with user input
if 'llm' in config_data and 'params' in config_data['llm']:
for param in LLM_CONF_OVERRIDE_KEY:
if param in kwargs and kwargs[param]:
logger.info(f"Overwrite with new {param} {kwargs[param]}")
config_data['llm']['params'][param] = kwargs[param]
# Create tasks for llm and each individual plugin
llm_config = config_data.get("llm", {})
plugin_configs = config_data.get('plugins', [])
# Create tasks for llm and each individual plugin
llm_task = asyncio.create_task(cls._async_init_llm(llm_config))
plugin_tasks = [asyncio.create_task(cls._async_init_plugin(plugin_config)) for
plugin_config in plugin_configs]
# Gather results
llm, *plugins = await asyncio.gather(llm_task, *plugin_tasks)
agent_instance.llm = llm
for plugin in plugins:
plugin_name, plugin_instance = plugin
agent_instance.add_plugin(plugin_name, plugin_instance)
return agent_instance
@classmethod
async def _async_init_llm(cls, llm_config):
llm_model_name = llm_config.get("module_name", None)
llm_class_name = llm_config.get("class_name", None)
if not llm_model_name or not llm_class_name:
raise InputErrorException("Agent LLM module_name and class_name required, please check your config")
module = import_module(llm_model_name)
clazz = getattr(module, llm_class_name)
assert issubclass(clazz, BaseLLM), f"{clazz} is not a subclass of BaseLLM"
llm_instance = await clazz.create(config_data=llm_config)
return llm_instance
@classmethod
async def _async_init_plugin(cls, plugin_config):
if plugin_config.get('type', "") == 'agent':
# Agent as plugin
agent = await BaseAgent.async_from_config_path_and_kwargs(plugin_config['config'])
return plugin_config['name'], agent
else:
# Tool as plugin
params = plugin_config.get('params', dict())
name = plugin_config.get('name', None)
config = plugin_config['config']
tool = await BaseTool.async_from_config(config_input=config, **params)
if name is None:
name = tool.name
logger.info("Init tool with name [{}], and description [{}]".format(name, tool.description))
return name, tool
@abstractmethod
def run(self, *args, **kwargs) -> [AgentResponse, None]:
"""Abstract method to be overridden by child classes for running the agent.
:return: The output of the agent.
:rtype: AgentOutput
"""
pass
async def async_run(self, *args, **kwargs) -> AsyncGenerator[AgentResponse, None]:
"""Abstract method to be overridden by child classes for running the agent.
:return: The output of the agent.
"""
yield self.run(*args, **kwargs)
def _get_plugin_function_map(self, method_name: str) -> Dict[str, Callable]:
if method_name == "run" and self.__plugin_tool_function:
return self.__plugin_tool_function
elif method_name == "async_run" and self.__plugin_tool_async_function:
return self.__plugin_tool_async_function
function_map = {}
for name, plugin_tool in self.plugins_map.items():
if isinstance(plugin_tool, (BaseTool, BaseAgent)):
function_map[name] = getattr(plugin_tool, method_name)
else:
logger.warning(f"No support for plugin name {name} of type {type(plugin_tool)}")
if method_name == "run":
self.__plugin_tool_function = function_map
elif method_name == "async_run":
self.__plugin_tool_async_function = function_map
return function_map
def get_plugin_tool_function(self) -> Dict[str, Callable]:
"""Format the function map for the function API.
:return: The function map.
:rtype: Dict[str, Callable]
"""
return self._get_plugin_function_map("run")
def get_plugin_tool_async_function(self) -> Dict[str, Callable]:
"""Format the function map for the function API.
:return: The function map.
:rtype: Dict[str, Callable]
"""
return self._get_plugin_function_map("async_run")
def _get_plugin_description(self):
if self.__plugin_tool_description:
return self.__plugin_tool_description
descriptions = ""
try:
for plugin_name, plugin in self.plugins_map.items():
descriptions += f"{plugin_name}[input]: {plugin.description}\n"
except Exception as e:
err_msg = "Failed to get plugin tool name and description. error: {}".format(str(e))
raise InputErrorException(err_msg) from e
self.__plugin_tool_description = descriptions
return descriptions
def clear(self):
"""
Clear and reset the agent.
"""
pass
|